|
| 1 | +# ------------------------------------------------------------------------- |
| 2 | +# Copyright (c) Microsoft Corporation. All rights reserved. |
| 3 | +# Licensed under the MIT License. |
| 4 | +# -------------------------------------------------------------------------- |
| 5 | +import argparse |
| 6 | +import tempfile |
| 7 | +import unittest |
| 8 | +from pathlib import Path |
| 9 | + |
| 10 | +from tokenizers import Tokenizer |
| 11 | +from tokenizers.models import WordLevel |
| 12 | +from tokenizers.pre_tokenizers import Whitespace |
| 13 | + |
| 14 | +DEFAULT_WHISPER_MODEL_IDS = ( |
| 15 | + "openai/whisper-tiny", |
| 16 | + "microsoft/whisper-base", |
| 17 | +) |
| 18 | +MAX_ARTIFACT_SIZE_BYTES = 1024 * 1024 |
| 19 | + |
| 20 | + |
| 21 | +def _run_cli_main(args): |
| 22 | + from olive.cli.launcher import main as cli_main |
| 23 | + |
| 24 | + cli_main(args) |
| 25 | + |
| 26 | + |
| 27 | +def _save_local_tiny_whisper(model_path: Path): |
| 28 | + from transformers import PreTrainedTokenizerFast, WhisperConfig, WhisperForConditionalGeneration |
| 29 | + |
| 30 | + model = WhisperForConditionalGeneration( |
| 31 | + WhisperConfig.from_dict( |
| 32 | + { |
| 33 | + "num_mel_bins": 80, |
| 34 | + "encoder_layers": 2, |
| 35 | + "encoder_attention_heads": 4, |
| 36 | + "decoder_layers": 2, |
| 37 | + "decoder_attention_heads": 4, |
| 38 | + "d_model": 64, |
| 39 | + "encoder_ffn_dim": 128, |
| 40 | + "decoder_ffn_dim": 128, |
| 41 | + "max_source_positions": 16, |
| 42 | + "max_target_positions": 16, |
| 43 | + "pad_token_id": 0, |
| 44 | + "bos_token_id": 1, |
| 45 | + "eos_token_id": 2, |
| 46 | + "decoder_start_token_id": 1, |
| 47 | + } |
| 48 | + ) |
| 49 | + ) |
| 50 | + model.save_pretrained(model_path) |
| 51 | + |
| 52 | + tokenizer = Tokenizer( |
| 53 | + WordLevel( |
| 54 | + vocab={"<pad>": 0, "<bos>": 1, "<eos>": 2, "hello": 3, "world": 4}, |
| 55 | + unk_token="<pad>", |
| 56 | + ) |
| 57 | + ) |
| 58 | + tokenizer.pre_tokenizer = Whitespace() |
| 59 | + PreTrainedTokenizerFast( |
| 60 | + tokenizer_object=tokenizer, |
| 61 | + bos_token="<bos>", |
| 62 | + eos_token="<eos>", |
| 63 | + pad_token="<pad>", |
| 64 | + ).save_pretrained(model_path) |
| 65 | + |
| 66 | + |
| 67 | +def _run_whisper_capture_onnx_flow(tmp_path: Path, model_id: str, precision: str = "fp32"): |
| 68 | + """Export a tiny local Whisper model to ONNX via capture-onnx-graph with Model Builder.""" |
| 69 | + model_name = model_id.replace("/", "--") |
| 70 | + model_path = tmp_path / "models" / model_name |
| 71 | + run_output_dir = tmp_path / f"{model_name}-onnx" |
| 72 | + |
| 73 | + _save_local_tiny_whisper(model_path) |
| 74 | + _run_cli_main( |
| 75 | + [ |
| 76 | + "capture-onnx-graph", |
| 77 | + "-m", |
| 78 | + str(model_path), |
| 79 | + "--use_model_builder", |
| 80 | + "--precision", |
| 81 | + precision, |
| 82 | + "--output_path", |
| 83 | + str(run_output_dir), |
| 84 | + ] |
| 85 | + ) |
| 86 | + |
| 87 | + return run_output_dir |
| 88 | + |
| 89 | + |
| 90 | +class TestCliWhisperSmoke(unittest.TestCase): |
| 91 | + """Smoke tests for Whisper encoder-decoder models exported via capture-onnx-graph.""" |
| 92 | + |
| 93 | + model_ids = DEFAULT_WHISPER_MODEL_IDS |
| 94 | + precision = "fp32" |
| 95 | + workdir = None |
| 96 | + |
| 97 | + def test_whisper_capture_onnx_graph(self): |
| 98 | + """Verify that Whisper encoder and decoder are exported to ONNX successfully.""" |
| 99 | + if self.workdir is None: |
| 100 | + with tempfile.TemporaryDirectory() as temp_dir: |
| 101 | + self._assert_whisper_export(Path(temp_dir)) |
| 102 | + else: |
| 103 | + workdir = Path(self.workdir) |
| 104 | + workdir.mkdir(parents=True, exist_ok=True) |
| 105 | + self._assert_whisper_export(workdir) |
| 106 | + |
| 107 | + def _assert_whisper_export(self, tmp_path: Path): |
| 108 | + expected_encoder_file = "encoder.onnx" |
| 109 | + expected_decoder_file = "decoder.onnx" |
| 110 | + for model_id in self.model_ids: |
| 111 | + with self.subTest(model_id=model_id): |
| 112 | + run_output_dir = _run_whisper_capture_onnx_flow(tmp_path, model_id, self.precision) |
| 113 | + output_files = self._list_relative_files(run_output_dir) |
| 114 | + assert expected_encoder_file in output_files, ( |
| 115 | + f"Expected {expected_encoder_file} in output, got: {output_files}" |
| 116 | + ) |
| 117 | + assert expected_decoder_file in output_files, ( |
| 118 | + f"Expected {expected_decoder_file} in output, got: {output_files}" |
| 119 | + ) |
| 120 | + self._assert_file_size_below_limit(run_output_dir / expected_encoder_file) |
| 121 | + self._assert_file_size_below_limit(run_output_dir / expected_decoder_file) |
| 122 | + |
| 123 | + def _assert_file_size_below_limit(self, path: Path): |
| 124 | + assert path.exists() |
| 125 | + assert path.stat().st_size < MAX_ARTIFACT_SIZE_BYTES |
| 126 | + |
| 127 | + @staticmethod |
| 128 | + def _list_relative_files(path: Path): |
| 129 | + return {file_path.relative_to(path).as_posix() for file_path in path.rglob("*") if file_path.is_file()} |
| 130 | + |
| 131 | + |
| 132 | +def _parse_args(): |
| 133 | + parser = argparse.ArgumentParser(add_help=False) |
| 134 | + parser.add_argument("--workdir") |
| 135 | + parser.add_argument("--model-id", dest="model_ids", action="append") |
| 136 | + return parser.parse_known_args() |
| 137 | + |
| 138 | + |
| 139 | +if __name__ == "__main__": |
| 140 | + parsed_args, remaining = _parse_args() |
| 141 | + if parsed_args.workdir: |
| 142 | + TestCliWhisperSmoke.workdir = Path(parsed_args.workdir) |
| 143 | + if parsed_args.model_ids: |
| 144 | + TestCliWhisperSmoke.model_ids = tuple(parsed_args.model_ids) |
| 145 | + unittest.main(argv=[__file__, *remaining]) |
0 commit comments