diff --git a/src/transformers/models/whisper/generation_whisper.py b/src/transformers/models/whisper/generation_whisper.py index 93dd29ae0fed..4a7c45ce3a8b 100644 --- a/src/transformers/models/whisper/generation_whisper.py +++ b/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] diff --git a/src/transformers/pipelines/automatic_speech_recognition.py b/src/transformers/pipelines/automatic_speech_recognition.py index b29f4160786d..2470849b4895 100644 --- a/src/transformers/pipelines/automatic_speech_recognition.py +++ b/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: diff --git a/tests/pipelines/test_pipelines_automatic_speech_recognition.py b/tests/pipelines/test_pipelines_automatic_speech_recognition.py index 0e3f2246cc57..0870cb934795 100644 --- a/tests/pipelines/test_pipelines_automatic_speech_recognition.py +++ b/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