Fix whisper return_language with return_timestamp=word

src/transformers/models/whisper/generation_whisper.py

@@ -411,6 +411,7 @@ def generate(
         return_dict_in_generate: Optional[bool] = None,
         force_unique_generate_call: Optional[bool] = None,
         monitor_progress: Optional[Callable[[torch.Tensor], None]] = None,
+        keep_special_tokens: Optional[int] = 0,
         **kwargs,
     ):
         """
@@ -540,6 +541,9 @@ def generate(
                 takes a tensor argument `p` of shape `(n, 2)`, where `n` is the batch size. `p[i, 0]`  contains the
                 index of the audio frame that is currently being transcribed for batch item `i`. `p[i, 1]` contains
                 the total number of frames for batch item `i`. No return value is expected.
+            keep_special_tokens (`int`, *optional*, defaults to 0):
+                The number of special tokens to keep at the beginning of the generated sequence.
+                Defaults to 0. -1 means keep all special tokens.
             kwargs (`dict[str, Any]`, *optional*):
                 Ad hoc parametrization of `generate_config` and/or additional model-specific kwargs that will be
                 forwarded to the `forward` function of the model. If the model is an encoder-decoder model, encoder
@@ -880,6 +884,7 @@ def generate(
                 is_shortform=is_shortform,
                 batch_size=batch_size,
                 attention_mask=attention_mask,
+                keep_special_tokens=keep_special_tokens,
                 kwargs=kwargs,
             )
 
@@ -904,6 +909,7 @@ def generate(
                     idx=i,
                     return_token_timestamps=return_token_timestamps,
                     decoder_input_ids=decoder_input_ids,
+                    num_special_tokens=keep_special_tokens,
                 )
 
                 seek[prev_i] += segment_offset
@@ -996,6 +1002,7 @@ def generate_with_fallback(
         is_shortform,
         batch_size,
         attention_mask,
+        keep_special_tokens,
         kwargs,
     ):
         kwargs = copy.copy(kwargs)
@@ -1058,6 +1065,7 @@ def generate_with_fallback(
                 is_shortform=is_shortform,
                 seek=seek,
                 batch_idx_map=batch_idx_map,
+                keep_special_tokens=keep_special_tokens,
             )
 
             if cur_bsz < batch_size:
@@ -1146,13 +1154,17 @@ def _postprocess_outputs(
         is_shortform,
         seek,
         batch_idx_map,
+        keep_special_tokens,
     ):
         # remove all previously passed decoder input ids
         # should happen only if it is the first generated segment
         start_idx = decoder_input_ids.shape[-1]
 
+        if keep_special_tokens == -1:
+            keep_special_tokens = start_idx
+
         if isinstance(seek_outputs, torch.Tensor):
-            return seek_outputs[:, start_idx:], seek_outputs
+            return seek_outputs[:, start_idx - keep_special_tokens :], seek_outputs
 
         if return_token_timestamps and hasattr(generation_config, "alignment_heads"):
             num_frames = getattr(generation_config, "num_frames")
@@ -1201,7 +1213,7 @@ def split_by_batch_index(values, key, batch_idx, is_shortform, beam_indices=None
 
             return values[batch_idx].cpu()
 
-        sequence_tokens = seek_outputs["sequences"][:, start_idx:]
+        sequence_tokens = seek_outputs["sequences"][:, start_idx - keep_special_tokens :]
         seek_outputs = [
             {
                 k: split_by_batch_index(v, k, i, is_shortform, beam_indices=seek_outputs.get("beam_indices"))
@@ -2009,6 +2021,7 @@ def _retrieve_segment(
         idx,
         return_token_timestamps,
         decoder_input_ids,
+        num_special_tokens,
     ):
         # find the predicted "end of segment" predictions of Whisper
         # "end of segment" predictions occur whenever Whisper predicts a timestamp token
@@ -2020,6 +2033,9 @@ def _retrieve_segment(
         idx_offset = decoder_input_ids.shape[-1]
         device = seek_sequence.device
 
+        if num_special_tokens == -1:
+            num_special_tokens = idx_offset
+
         # If whisper predicted a "end of segment" via a timestep token, let's go ever each
         # "end of segment" prediction and slice the decoding into segments accordingly
         if len(timestamp_segment_indices) > 0:
@@ -2037,7 +2053,17 @@ def _retrieve_segment(
             for i, current_slice in enumerate(slices):
                 is_last_slice = i == len(slices) - 1
                 sliced_tokens = seek_sequence[last_slice:current_slice]
-                start_timestamp_pos = sliced_tokens[0] - timestamp_begin
+
+                start_timestamp_pos = None
+                for token in sliced_tokens:
+                    if token >= timestamp_begin:
+                        start_timestamp_pos = token - timestamp_begin
+                        break
+
+                if start_timestamp_pos is None:
+                    # This should not be possible. Fallback to previous logic.
+                    start_timestamp_pos = sliced_tokens[0] - timestamp_begin
+
                 idx_sliced_tokens = -1 if not is_last_slice or single_timestamp_ending else -2
                 end_timestamp_pos = sliced_tokens[idx_sliced_tokens] - timestamp_begin
                 segments.append(
@@ -2049,13 +2075,21 @@ def _retrieve_segment(
                         + end_timestamp_pos.to(torch.float32 if device.type == "mps" else torch.float64)
                         * time_precision,
                         "tokens": sliced_tokens,
-                        "idxs": (idx_offset + last_slice, idx_offset + current_slice),
+                        "idxs": (
+                            idx_offset + last_slice - num_special_tokens,
+                            idx_offset + current_slice - num_special_tokens,
+                        ),
                         "result": seek_outputs[idx],
                     }
                 )
                 if return_token_timestamps:
                     segments[-1]["token_timestamps"] = (
-                        token_timestamps[idx_offset + last_slice : idx_offset + current_slice] + time_offset[prev_idx]
+                        token_timestamps[
+                            idx_offset + last_slice - num_special_tokens : idx_offset
+                            + current_slice
+                            - num_special_tokens
+                        ]
+                        + time_offset[prev_idx]
                     )
                 last_slice = current_slice
 
@@ -2083,13 +2117,16 @@ def _retrieve_segment(
                     "start": time_offset[prev_idx],
                     "end": time_offset[prev_idx] + last_timestamp_pos * time_precision,
                     "tokens": seek_sequence,
-                    "idxs": (idx_offset, idx_offset + len(seek_sequence)),
+                    "idxs": (idx_offset - num_special_tokens, idx_offset + len(seek_sequence) - num_special_tokens),
                     "result": seek_outputs[idx],
                 }
             ]
             if return_token_timestamps:
                 segments[-1]["token_timestamps"] = (
-                    token_timestamps[idx_offset : idx_offset + len(seek_sequence)] + time_offset[prev_idx]
+                    token_timestamps[
+                        idx_offset - num_special_tokens : idx_offset + len(seek_sequence) - num_special_tokens
+                    ]
+                    + time_offset[prev_idx]
                 )
             segment_offset = seek_num_frames[prev_idx]
 

src/transformers/pipelines/automatic_speech_recognition.py

@@ -343,6 +343,7 @@ def _sanitize_parameters(
         if return_language is not None:
             if self.type != "seq2seq_whisper":
                 raise ValueError("Only Whisper can return language for now.")
+            forward_params["return_language"] = return_language
             postprocess_params["return_language"] = return_language
 
         if getattr(self, "assistant_model", None) is not None:
@@ -499,7 +500,7 @@ def preprocess(self, inputs, chunk_length_s=0, stride_length_s=None):
                 processed["stride"] = stride
             yield {"is_last": True, **processed, **extra}
 
-    def _forward(self, model_inputs, return_timestamps=False, **generate_kwargs):
+    def _forward(self, model_inputs, return_timestamps=False, return_language=False, **generate_kwargs):
         attention_mask = model_inputs.pop("attention_mask", None)
         stride = model_inputs.pop("stride", None)
         num_frames = model_inputs.pop("num_frames", None)
@@ -528,6 +529,10 @@ def _forward(self, model_inputs, return_timestamps=False, **generate_kwargs):
                 if return_timestamps == "word":
                     generate_kwargs["return_token_timestamps"] = True
                     generate_kwargs["return_segments"] = True
+                    if return_language:
+                        # The First three special tokens will be <|startoftranscript|><|language|><|task(transcribe/translate)|>
+                        # Here we ask for two tokens to be preserved so <|language|> is returned.
+                        generate_kwargs["keep_special_tokens"] = 2
 
             # User-defined `generation_config` passed to the pipeline call take precedence
             if "generation_config" not in generate_kwargs:

tests/pipelines/test_pipelines_automatic_speech_recognition.py

@@ -418,59 +418,72 @@ def test_return_timestamps_in_preprocess(self):
 
     @slow
     @require_torch
-    @unittest.skip("TODO (joao, eustache): this test is failing, find the breaking PR and fix the cause or the test")
-    def test_return_timestamps_and_language_in_preprocess(self):
-        pipe = pipeline(
-            task="automatic-speech-recognition",
-            model="openai/whisper-tiny",
-            chunk_length_s=8,
-            stride_length_s=1,
-            return_language=True,
+    def test_return_timestamps_and_language(self):
+        pipe = pipeline(task="automatic-speech-recognition", model="openai/whisper-tiny")
+        data = load_dataset(
+            "openslr/librispeech_asr",
+            "clean",
+            split="test",
+            streaming=True,
         )
-        data = load_dataset("openslr/librispeech_asr", "clean", split="test", streaming=True)
         sample = next(iter(data))
 
         res = pipe(sample["audio"]["array"])
         self.assertEqual(
             res,
-            {
-                "text": " Conquered returned to its place amidst the tents.",
-                "chunks": [{"language": "english", "text": " Conquered returned to its place amidst the tents."}],
-            },
+            {"text": " Concord returned to its place amidst the tents."},
         )
 
         res = pipe(sample["audio"]["array"], return_timestamps=True)
         self.assertEqual(
             res,
             {
-                "text": " Conquered returned to its place amidst the tents.",
+                "text": " Concord returned to its place amidst the tents.",
                 "chunks": [
                     {
                         "timestamp": (0.0, 3.36),
-                        "language": "english",
-                        "text": " Conquered returned to its place amidst the tents.",
+                        "text": " Concord returned to its place amidst the tents.",
                     }
                 ],
             },
         )
 
         res = pipe(sample["audio"]["array"], return_timestamps="word")
+        self.assertEqual(
+            res,
+            {
+                "text": " Concord returned to its place amidst the tents.",
+                "chunks": [
+                    {"text": " Concord", "timestamp": (0.0, 1.2)},
+                    {"text": " returned", "timestamp": (1.2, 1.62)},
+                    {"text": " to", "timestamp": (1.62, 1.86)},
+                    {"text": " its", "timestamp": (1.86, 2.02)},
+                    {"text": " place", "timestamp": (2.02, 2.28)},
+                    {"text": " amidst", "timestamp": (2.28, 2.82)},
+                    {"text": " the", "timestamp": (2.82, 2.98)},
+                    {"text": " tents.", "timestamp": (2.98, 3.48)},
+                ],
+            },
+        )
+
+        res = pipe(sample["audio"]["array"], return_language=True, return_timestamps="word")
         # fmt: off
         self.assertEqual(
             res,
             {
-                'text': ' Conquered returned to its place amidst the tents.',
-                'chunks': [
-                    {"language": "english",'text': ' Conquered', 'timestamp': (0.5, 1.2)},
-                    {"language": "english", 'text': ' returned', 'timestamp': (1.2, 1.64)},
-                    {"language": "english",'text': ' to', 'timestamp': (1.64, 1.84)},
-                    {"language": "english",'text': ' its', 'timestamp': (1.84, 2.02)},
-                    {"language": "english",'text': ' place', 'timestamp': (2.02, 2.28)},
-                    {"language": "english",'text': ' amidst', 'timestamp': (2.28, 2.8)},
-                    {"language": "english",'text': ' the', 'timestamp': (2.8, 2.98)},
-                    {"language": "english",'text': ' tents.', 'timestamp': (2.98, 3.48)},
+                "text": " Concord returned to its place amidst the tents.",
+                "chunks": [
+                    {"text": " Concord","timestamp": (0.0, 1.2),"language": "english",},
+                    {"text": " returned","timestamp": (1.2, 1.62),"language": "english",},
+                    {"text": " to", "timestamp": (1.62, 1.86), "language": "english"},
+                    {"text": " its", "timestamp": (1.86, 2.02), "language": "english"},
+                    {"text": " place","timestamp": (2.02, 2.28),"language": "english",},
+                    {"text": " amidst","timestamp": (2.28, 2.82),"language": "english",},
+                    {"text": " the", "timestamp": (2.82, 2.98), "language": "english"},
+                    {"text": " tents.","timestamp": (2.98, 3.48),"language": "english",},
                 ],
             },
+
         )
         # fmt: on