-
Notifications
You must be signed in to change notification settings - Fork 3k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Add timestamp logits processor for whisper (#15853)
Enable timestamp estimation and logits processing for Whisper model.
- Loading branch information
Showing
10 changed files
with
254 additions
and
7 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
75 changes: 75 additions & 0 deletions
75
onnxruntime/test/python/transformers/test_whisper_timestamp_processor.py
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,75 @@ | ||
| # ------------------------------------------------------------------------- | ||
| # Copyright (c) Microsoft Corporation. All rights reserved. | ||
| # Licensed under the MIT License. See License.txt in the project root for | ||
| # license information. | ||
| # -------------------------------------------------------------------------- | ||
|
|
||
| import unittest | ||
|
|
||
| import numpy as np | ||
| import pytest | ||
| import torch | ||
|
|
||
| from onnxruntime import InferenceSession, SessionOptions | ||
|
|
||
|
|
||
| class TestTimestampProcessor(unittest.TestCase): | ||
| def generate_model(self, arguments: str): | ||
| from onnxruntime.transformers.models.whisper.convert_to_onnx import main as whisper_to_onnx | ||
|
|
||
| whisper_to_onnx(arguments.split()) | ||
|
|
||
| def generate_dataset(self): | ||
| from datasets import load_dataset | ||
| from transformers import AutoProcessor | ||
|
|
||
| processor = AutoProcessor.from_pretrained("openai/whisper-tiny") | ||
| ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation") | ||
| inputs = processor(ds[0]["audio"]["array"], return_tensors="pt") | ||
| input_features = inputs.input_features | ||
| return [input_features, processor] | ||
|
|
||
| def run_timestamp(self, provider: str): | ||
| self.generate_model("-m openai/whisper-tiny --optimize_onnx --precision fp32 -l -e") | ||
| [input_features, processor] = self.generate_dataset() | ||
| model_path = "./onnx_models/openai/whisper-tiny_beamsearch.onnx" | ||
| sess_options = SessionOptions() | ||
| sess_options.log_severity_level = 4 | ||
| sess = InferenceSession(model_path, sess_options, providers=[provider]) | ||
| input_data = input_features.repeat(1, 1, 1) | ||
| ort_inputs = { | ||
| "input_features": np.float32(input_data.cpu().numpy()), | ||
| "max_length": np.array([128], dtype=np.int32), | ||
| "min_length": np.array([0], dtype=np.int32), | ||
| "num_beams": np.array([1], dtype=np.int32), | ||
| "num_return_sequences": np.array([1], dtype=np.int32), | ||
| "length_penalty": np.array([1.0], dtype=np.float32), | ||
| "repetition_penalty": np.array([1.0], dtype=np.float32), | ||
| "logits_processor": np.array([1], dtype=np.int32), | ||
| } | ||
| ort_out = sess.run(None, ort_inputs) | ||
| ort_out_tensor = torch.from_numpy(ort_out[0]) | ||
| ort_transcription = processor.batch_decode( | ||
| ort_out_tensor[0][0].view(1, -1), skip_special_tokens=True, output_offsets=True | ||
| ) | ||
| expected_transcription = [ | ||
| { | ||
| "text": " Mr. Quilter is the apostle of the middle classes and we are glad to welcome his gospel.", | ||
| "offsets": [ | ||
| { | ||
| "text": " Mr. Quilter is the apostle of the middle classes and we are glad to welcome his gospel.", | ||
| "timestamp": (0.0, 5.44), | ||
| } | ||
| ], | ||
| } | ||
| ] | ||
| self.assertEqual(ort_transcription, expected_transcription) | ||
|
|
||
| @pytest.mark.slow | ||
| def test_timestamp_cpu(self): | ||
| provider = "CPUExecutionProvider" | ||
| self.run_timestamp(provider) | ||
|
|
||
|
|
||
| if __name__ == "__main__": | ||
| unittest.main() |