The Whisper large-v3 model exported to ONNX does not return the end timestamp for the last chunk

[IlyaPikin]

System Info

optimum: 1.19.1
onnxruntime: 1.17.3
onnx: 1.16.0
python: 3.11
transformers: 4.37.1
platform: Ubuntu 20.04.6 LTS

Who can help?

@JingyaHuang @echarlaix @michaelbenayoun

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• My own task or dataset (give details below)

Reproduction

Running the transformers pipeline for the original whisper-large-v3 model returns the correct timestamps of all chunks for files of any duration. For example, for a file with a duration of 33 seconds, the code below gives the following result:

import torch
from transformers import WhisperForConditionalGeneration, WhisperProcessor, GenerationConfig, pipeline


model_name = "openai/whisper-large-v3"
processor = WhisperProcessor.from_pretrained(model_name)
model = WhisperForConditionalGeneration.from_pretrained(model_name,
                                                torch_dtype=torch.float32,
                                                low_cpu_mem_usage=False,
                                                use_safetensors=False)
model.eval()

generation_config = GenerationConfig.from_pretrained(model_name)
model.generation_config = generation_config
model.to("cpu")

pipe = pipeline(
    "automatic-speech-recognition",
    model=model,
    tokenizer=processor.tokenizer,
    feature_extractor=processor.feature_extractor,
    chunk_length_s=30,
    batch_size=16,
    device="cpu",
)

audio_path = "30.192.wav"

transcript = pipe(audio_path, return_timestamps=True, batch_size=16)
print(transcript)

Output

{'text': ' Глаза Степана Аркадьевича весело заблестели, и он задумался, улыбаясь. Да, хорошо было, очень хорошо. Много еще что-то там было отличного. Да не скажешь словами, и мыслями даже наяву не выразишь. И, заметив полосу света, пробившую сбоку одной из суконных стор, он весело скинул ноги с дивана, отыскал ими, шитые женой, подарок о дню рождения в прошлом году, обделанный в золотистые Софьян туфли, и по старой девятилетней привычке, не вставая, потянулся рукой к тому месту, где в спальне у него висел халат. И тут он вспомнил вдруг, как и почему он спит не в спальне жены, а в кабинете. Улыбка исчезла с его лица, он сморщил лоб.', 'chunks': [{'timestamp': (0.0, 3.52), 'text': ' Глаза Степана Аркадьевича весело заблестели, и он задумался, улыбаясь.'}, {'timestamp': (3.66, 5.34), 'text': ' Да, хорошо было, очень хорошо.'}, {'timestamp': (5.9, 7.2), 'text': ' Много еще что-то там было отличного.'}, {'timestamp': (7.26, 9.56), 'text': ' Да не скажешь словами, и мыслями даже наяву не выразишь.'}, {'timestamp': (9.92, 13.4), 'text': ' И, заметив полосу света, пробившую сбоку одной из суконных стор,'}, {'timestamp': (13.72, 15.46), 'text': ' он весело скинул ноги с дивана,'}, {'timestamp': (15.54, 18.58), 'text': ' отыскал ими, шитые женой, подарок о дню рождения в прошлом году,'}, {'timestamp': (19.1, 20.82), 'text': ' обделанный в золотистые Софьян туфли,'}, {'timestamp': (21.02, 23.34), 'text': ' и по старой девятилетней привычке, не вставая,'}, {'timestamp': (23.68, 26.26), 'text': ' потянулся рукой к тому месту, где в спальне у него висел халат.'}, {'timestamp': (26.68, 30.0), 'text': ' И тут он вспомнил вдруг, как и почему он спит не в спальне жены, а в кабинете.'}, {'timestamp': (30.42, 32.52), 'text': ' Улыбка исчезла с его лица, он сморщил лоб.'}]}

After converting to onnx using this command:

optimum-cli export onnx --model openai/whisper-large-v3 whisper-large-v3-with-past/ --task automatic-speech-recognition-with-past --opset 14 --framework pt --device cpu --atol 0.01 --no-post-process

and running the equivalent code:

import os
from transformers import WhisperProcessor, PretrainedConfig, GenerationConfig, pipeline
from optimum.onnxruntime import ORTModelForSpeechSeq2Seq

model_path = "../tools/quantization/whisper-large-v3-with-past"

processor = WhisperProcessor.from_pretrained(model_path)
model_config = PretrainedConfig.from_pretrained(model_path)

sessions = ORTModelForSpeechSeq2Seq.load_model(
    os.path.join(model_path, 'encoder_model.onnx'),
    os.path.join(model_path, 'decoder_model.onnx'),
    os.path.join(model_path, 'decoder_with_past_model.onnx'))
model = ORTModelForSpeechSeq2Seq(sessions[0], sessions[1], model_config, 
                                    model_path, sessions[2])

generation_config = GenerationConfig.from_pretrained(model_path)
model.generation_config = generation_config
model.to("cpu")

pipe = pipeline(
    "automatic-speech-recognition",
    model=model,
    tokenizer=processor.tokenizer,
    feature_extractor=processor.feature_extractor,
    chunk_length_s=30,
    batch_size=16,
    device="cpu",
)

audio_path = "30.192.wav"

transcript = pipe(audio_path, return_timestamps=True, batch_size=16)
print(transcript)

The transcript contains None instead of the end timestamp for the last chunk, although there is no word cut off in the middle:
Output

{'text': ' Глаза Степана Аркадьевича весело заблестели, и он задумался, улыбаясь. Да, хорошо было, очень хорошо. Много еще что-то там было отличного, да не скажешь словами и мыслями даже наяву не выразишь. И, заметив полосу света, пробившую сбоку одной из суконных стор, он весело скинул ноги с дивана, отыскал ими, шитые женой, подарок о дню рождения в прошлом году, обделанный в золотистые Софьян туфли, и по старой девятилетней привычке, не вставая, потянулся рукой к тому месту, где в спальне у него висел халат. И тут он вспомнил вдруг, как и почему он спит не в спальне жены, а в кабинете. Улыбка исчезла с его лица, он сморщил лоб.', 'chunks': [{'timestamp': (0.0, 9.56), 'text': ' Глаза Степана Аркадьевича весело заблестели, и он задумался, улыбаясь. Да, хорошо было, очень хорошо. Много еще что-то там было отличного, да не скажешь словами и мыслями даже наяву не выразишь.'}, {'timestamp': (9.92, None), 'text': ' И, заметив полосу света, пробившую сбоку одной из суконных стор, он весело скинул ноги с дивана, отыскал ими, шитые женой, подарок о дню рождения в прошлом году, обделанный в золотистые Софьян туфли, и по старой девятилетней привычке, не вставая, потянулся рукой к тому месту, где в спальне у него висел халат. И тут он вспомнил вдруг, как и почему он спит не в спальне жены, а в кабинете. Улыбка исчезла с его лица, он сморщил лоб.'}]}

In the case of files with a duration of less than 30 seconds, the ONNX model does not return timestamps at all:

{'text': ' Обдумал я все это и вроде бы даже облегчение почувствовал, что в зону мне сегодня идти не надо. Только как это все поделикатнее сообщить Кириллу?', 'chunks': [{'timestamp': (None, None), 'text': ' Обдумал я все это и вроде бы даже облегчение почувствовал, что в зону мне сегодня идти не надо. Только как это все поделикатнее сообщить Кириллу?'}]}

Expected:

{'text': ' обдумал я все это и вроде бы даже облегчение почувствовал, что в зону мне сегодня идти не надо. Только как это все поделикатнее сообщить Кириллу?', 'chunks': [{'timestamp': (0.0, 3.92), 'text': ' обдумал я все это и вроде бы даже облегчение почувствовал,'}, {'timestamp': (4.02, 6.54), 'text': ' что в зону мне сегодня идти не надо.'}, {'timestamp': (6.54, 9.86), 'text': ' Только как это все поделикатнее сообщить Кириллу?'}]}

So now there are 2 problems:

1. If the audio is less than 30 seconds, there are no timestamps at all.
2. If it is longer, then there are no timestamps for the last chunk.

Audio files:
audio_files.zip

Expected behavior

The code which uses the ONNX model should work the same as the version using the pytorch model.