Why Gemfury?
Push, build, and install
RubyGems
npm packages
Python packages
Maven artifacts
PHP packages
Go Modules
Debian packages
RPM packages
NuGet packages
turingmotors
turingmotors / onnxruntime-gpu python
Repository URL to install this package:
|
Version:
1.23.2 ▾
|
# -------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License. See License.txt in the project root for
# license information.
# --------------------------------------------------------------------------
import json
import logging
import os
from pathlib import Path
import numpy as np
import torch
from convert_generation import add_cache_indirection_to_mha, add_output_qk_to_mha, fix_past_sequence_length
from optimizer import optimize_model
from transformers import AutoTokenizer, WhisperConfig, WhisperForConditionalGeneration, WhisperProcessor
from whisper_decoder import WhisperDecoder
from whisper_encoder import WhisperEncoder
from whisper_encoder_decoder_init import WhisperEncoderDecoderInit
from whisper_jump_times import WhisperJumpTimes
from onnxruntime import InferenceSession
logger = logging.getLogger(__name__)
PRETRAINED_WHISPER_MODELS = [
"whisper-tiny",
"whisper-tiny.en",
"whisper-base",
"whisper-base.en",
"whisper-small",
"whisper-small.en",
"whisper-medium",
"whisper-medium.en",
"whisper-large",
"whisper-large-v2",
"whisper-large-v3",
"whisper-large-v3-turbo",
]
class WhisperHelper:
@staticmethod
def get_onnx_path(
output_dir: str,
model_name_or_path: str,
suffix: str = "",
new_folder: bool = False,
) -> str:
"""Build onnx path
Args:
output_dir (str): output directory
model_name_or_path (str): pretrained model name, or path to the model checkpoint
suffix (str, optional): suffix like "_encoder" or "_decoder_fp16" will be appended to file name. Defaults to None.
new_folder (bool, optional): create a new directory for the model. Defaults to False.
Returns:
str: path of onnx model
"""
model_name = model_name_or_path
if os.path.isdir(model_name_or_path):
model_name = Path(model_name_or_path).parts[-1]
else:
model_name = model_name.split("/")[-1]
model_name += suffix
directory = os.path.join(output_dir, model_name) if new_folder else output_dir
return os.path.join(directory, model_name + ".onnx")
@staticmethod
def save_processing(
model_name_or_path: str,
provider: str,
separate_encoder_and_decoder_init: bool,
use_decoder_masked_mha: bool,
output_qk: bool,
encoder_path: str,
decoder_path: str,
output_dir: str,
cache_dir: str,
) -> None:
config = WhisperConfig.from_pretrained(model_name_or_path, cache_dir=cache_dir)
config.save_pretrained(output_dir)
tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, cache_dir=cache_dir)
tokenizer.save_pretrained(output_dir)
processor = WhisperProcessor.from_pretrained(model_name_or_path, cache_dir=cache_dir)
processor.save_pretrained(output_dir)
# Return early since the next files are for ONNX Runtime GenAI
if separate_encoder_and_decoder_init:
return
audio_processor_cfg = {
"feature_extraction": {
"sequence": [
{"operation": {"name": "audio_decoder", "type": "AudioDecoder"}},
{
"operation": {
"name": "STFT",
"type": "STFTNorm",
"attrs": {
"n_fft": 400,
"frame_length": 400,
"hop_length": 160,
"_comment": [
0.0,
0.0000616908073425293,
0.0002467334270477295,
0.0005550682544708252,
0.000986635684967041,
0.0015413463115692139,
0.0022190213203430176,
0.0030195116996765137,
0.003942638635635376,
0.004988163709640503,
0.006155818700790405,
0.007445335388183594,
0.008856385946273804,
0.010388582944869995,
0.012041628360748291,
0.013815045356750488,
0.01570841670036316,
0.01772129535675049,
0.019853144884109497,
0.022103488445281982,
0.02447172999382019,
0.026957333087921143,
0.029559612274169922,
0.03227800130844116,
0.03511175513267517,
0.03806024789810181,
0.0411226749420166,
0.044298380613327026,
0.04758647084236145,
0.05098623037338257,
0.05449673533439636,
0.058117181062698364,
0.06184667348861694,
0.0656842589378357,
0.06962898373603821,
0.07367992401123047,
0.0778360664844513,
0.08209633827209473,
0.08645972609519958,
0.09092515707015991,
0.09549149870872498,
0.10015767812728882,
0.10492250323295593,
0.1097848117351532,
0.11474338173866272,
0.11979702115058899,
0.12494447827339172,
0.13018447160720825,
0.1355157196521759,
0.14093685150146484,
0.1464466154575348,
0.15204361081123352,
0.1577264666557312,
0.16349375247955322,
0.16934409737586975,
0.1752760112285614,
0.18128803372383118,
0.18737870454788208,
0.19354650378227234,
0.1997898817062378,
0.20610737800598145,
0.21249738335609436,
0.21895831823349,
0.2254886031150818,
0.23208662867546082,
0.23875075578689575,
0.24547931551933289,
0.2522706985473633,
0.25912320613861084,
0.26603513956069946,
0.27300477027893066,
0.2800304591655731,
0.2871103882789612,
0.29424285888671875,
0.30142611265182495,
0.30865830183029175,
0.31593772768974304,
0.3232625722885132,
0.3306310474872589,
0.3380413055419922,
0.34549152851104736,
0.352979838848114,
0.3605044484138489,
0.3680635094642639,
0.37565508484840393,
0.38327735662460327,
0.3909284174442291,
0.39860638976097107,
0.4063093662261963,
0.41403549909591675,
0.42178282141685486,
0.4295494258403778,
0.43733343482017517,
0.44513291120529175,
0.45294591784477234,
0.46077051758766174,
0.46860480308532715,
0.4764467775821686,
0.4842946231365204,
0.492146372795105,
0.5,
0.5078536868095398,
0.515705406665802,
0.5235532522201538,
0.5313953161239624,
0.5392295718193054,
0.5470541715621948,
0.5548672080039978,
0.562666654586792,
0.5704506635665894,
0.5782172679901123,
0.5859646201133728,
0.5936906933784485,
0.6013936996459961,
0.609071671962738,
0.6167227625846863,
0.6243450045585632,
0.6319366097450256,
0.6394955515861511,
0.6470202207565308,
0.6545085310935974,
0.6619587540626526,
0.6693689823150635,
0.6767374277114868,
0.6840623021125793,
0.691341757774353,
0.6985740065574646,
0.7057572603225708,
0.7128896713256836,
0.719969630241394,
0.7269952893257141,
0.7339649796485901,
0.7408769130706787,
0.7477294206619263,
0.7545207738876343,
0.761249303817749,
0.7679134607315063,
0.774511456489563,
0.7810417413711548,
0.7875027060508728,
0.7938927412033081,
0.800210177898407,
0.8064535856246948,
0.8126214146614075,
0.8187121152877808,
0.8247240781784058,
0.8306560516357422,
0.8365063667297363,
0.8422735929489136,
0.8479564785957336,
0.8535534143447876,
0.8590631484985352,
0.8644843101501465,
0.8698155879974365,
0.8750555515289307,
0.8802030086517334,
0.8852566480636597,
0.8902152180671692,
0.8950775265693665,
0.899842381477356,
0.9045084714889526,
0.9090749025344849,
0.9135403037071228,
0.9179036617279053,
0.9221639633178711,
0.9263200759887695,
0.9303710460662842,
0.9343158006668091,
0.9381533861160278,
0.941882848739624,
0.945503294467926,
0.9490138292312622,
0.9524135589599609,
0.9557017087936401,
0.9588773250579834,
0.961939811706543,
0.9648882746696472,
0.9677220582962036,
0.9704403877258301,
0.9730427265167236,
0.9755282998085022,
0.9778965711593628,
0.9801468849182129,
0.9822787046432495,
0.9842916131019592,
0.9861849546432495,
0.9879584312438965,
0.9896113872528076,
0.9911436438560486,
0.9925546646118164,
0.9938441514968872,
0.9950118064880371,
0.996057391166687,
0.9969804883003235,
0.997780978679657,
0.9984586238861084,
0.999013364315033,
0.9994449615478516,
0.9997532367706299,
0.9999383091926575,
1,
0.9999383091926575,
0.9997532367706299,
0.9994449615478516,
0.999013364315033,
0.9984586238861084,
0.997780978679657,
0.9969804286956787,
0.9960573315620422,
0.9950118064880371,
0.9938441514968872,
0.9925546646118164,
0.9911435842514038,
0.9896113872528076,
0.9879583716392517,
0.9861849546432495,
0.9842915534973145,
0.9822787046432495,
0.9801468253135681,
0.9778964519500732,
0.9755282402038574,
0.9730426073074341,
0.9704403877258301,
0.9677219390869141,
0.9648882150650024,
0.9619396924972534,
0.9588772654533386,
0.9557015895843506,
0.9524134397506714,
0.9490137100219727,
0.9455032348632812,
0.9418827295303345,
0.9381532669067383,
0.9343156814575195,
0.9303709268569946,
0.9263200759887695,
0.9221639633178711,
0.9179036617279053,
0.913540244102478,
0.9090747833251953,
0.9045084714889526,
0.8998422622680664,
0.8950774669647217,
0.8902151584625244,
0.8852565884590149,
0.8802029490470886,
0.8750554919242859,
0.869815468788147,
0.8644842505455017,
0.8590630888938904,
0.853553295135498,
0.8479562997817993,
0.842273473739624,
0.836506187915802,
0.8306558728218079,
0.8247239589691162,
0.8187118768692017,
0.8126212358474731,
0.8064534664154053,
0.8002099990844727,
0.793892502784729,
0.7875025272369385,
0.7810416221618652,
0.7745113372802734,
0.767913281917572,
0.7612491846084595,
0.7545205950737,
0.7477291822433472,
0.7408767342567444,
0.7339648008346558,
0.7269951105117798,
0.7199694514274597,
0.7128894925117493,
0.7057570219039917,
0.6985738277435303,
0.6913415789604187,
0.684062123298645,
0.6767372488975525,
0.6693688035011292,
0.6619585752487183,
0.6545083522796631,
0.6470199823379517,
0.6394953727722168,
0.6319363117218018,
0.6243447661399841,
0.6167224645614624,
0.6090714335441589,
0.601393461227417,
0.5936904549598694,
0.5859643220901489,
0.5782170295715332,
0.5704504251480103,
0.5626664161682129,
0.5548669099807739,
0.5470539331436157,
0.5392293334007263,
0.5313950181007385,
0.5235530138015747,
0.5157051682472229,
0.507853627204895,
0.5,
0.4921463429927826,
0.484294593334198,
0.4764467477798462,
0.46860471367836,
0.4607704281806946,
0.4529458284378052,
0.4451328217983246,
0.437333345413208,
0.42954933643341064,
0.4217827320098877,
0.4140354096889496,
0.4063093066215515,
0.3986063003540039,
0.39092832803726196,
0.3832772672176361,
0.37565499544143677,
0.36806342005729675,
0.3605043888092041,
0.35297977924346924,
0.3454914391040802,
0.338041216135025,
0.33063095808029175,
0.3232625126838684,
0.3159376382827759,
0.3086581826210022,
0.3014259934425354,
0.2942427396774292,
0.28711026906967163,
0.2800303101539612,
0.2730046510696411,
0.2660350203514099,
0.2591230869293213,
0.25227057933807373,
0.24547919631004333,
0.2387506067752838,
0.23208650946617126,
0.22548848390579224,
0.21895819902420044,
0.2124972641468048,
0.2061072587966919,
0.19978976249694824,
0.1935463547706604,
0.18737855553627014,
0.18128788471221924,
0.17527586221694946,
0.1693439483642578,
0.16349363327026367,
0.15772631764411926,
0.15204349160194397,
0.14644649624824524,
0.1409367322921753,
0.13551557064056396,
0.1301843225955963,
0.12494435906410217,
0.11979690194129944,
0.11474326252937317,
0.10978469252586365,
0.10492238402366638,
0.10015755891799927,
0.09549137949943542,
0.09092503786087036,
0.08645960688591003,
0.08209621906280518,
0.07783591747283936,
0.07367980480194092,
0.06962886452674866,
0.06568413972854614,
0.06184655427932739,
0.0581170916557312,
0.0544966459274292,
0.05098611116409302,
0.04758638143539429,
0.044298261404037476,
0.04112258553504944,
0.038060128688812256,
0.03511166572570801,
0.03227788209915161,
0.02955952286720276,
0.02695724368095398,
0.024471670389175415,
0.02210339903831482,
0.01985308527946472,
0.017721205949783325,
0.015708357095718384,
0.0138150155544281,
0.012041598558425903,
0.010388582944869995,
0.008856356143951416,
0.007445335388183594,
0.006155818700790405,
0.004988163709640503,
0.003942638635635376,
0.0030195116996765137,
0.0022190213203430176,
0.0015413165092468262,
0.000986635684967041,
0.0005550682544708252,
0.0002467334270477295,
0.0000616908073425293,
],
},
}
},
{
"operation": {
"name": "log_mel_spectrogram",
"type": "LogMelSpectrum",
"attrs": {"chunk_size": 30, "hop_length": 160, "n_fft": 400, "n_mel": config.num_mel_bins},
}
},
]
}
}
audio_processor_json = json.dumps(audio_processor_cfg, indent=4)
with open(os.path.join(output_dir, "audio_processor_config.json"), "w") as f:
f.write(audio_processor_json)
provider_options = [] if "cpu" in provider else [{f"{provider}": {}}]
genai_config = {
"model": {
"bos_token_id": config.bos_token_id,
"context_length": config.max_length,
"decoder": {
"session_options": {
"log_id": "onnxruntime-genai",
"provider_options": provider_options,
},
"filename": os.path.basename(decoder_path),
"head_size": config.d_model // config.decoder_attention_heads,
"hidden_size": config.d_model,
"inputs": {
"input_ids": "input_ids",
"past_key_names": "past_key_self_%d",
"past_value_names": "past_value_self_%d",
"cross_past_key_names": "past_key_cross_%d",
"cross_past_value_names": "past_value_cross_%d",
},
"outputs": {
"logits": "logits",
"present_key_names": "present_key_self_%d",
"present_value_names": "present_value_self_%d",
},
"num_attention_heads": config.decoder_attention_heads,
"num_hidden_layers": config.decoder_layers,
"num_key_value_heads": config.decoder_attention_heads,
},
"encoder": {
"session_options": {
"log_id": "onnxruntime-genai",
"provider_options": provider_options,
},
"filename": os.path.basename(encoder_path),
"head_size": config.d_model // config.encoder_attention_heads,
"hidden_size": config.d_model,
"inputs": {"audio_features": "audio_features"},
"outputs": {
"encoder_hidden_states": "encoder_hidden_states",
"cross_present_key_names": "present_key_cross_%d",
"cross_present_value_names": "present_value_cross_%d",
},
"num_attention_heads": config.encoder_attention_heads,
"num_hidden_layers": config.encoder_layers,
"num_key_value_heads": config.encoder_attention_heads,
},
"eos_token_id": config.eos_token_id,
"pad_token_id": config.pad_token_id,
"type": "whisper",
"vocab_size": config.vocab_size,
},
"search": {
"diversity_penalty": 0.0,
"do_sample": False,
"early_stopping": True,
"length_penalty": 1.0,
"max_length": config.max_length,
"min_length": 0,
"no_repeat_ngram_size": 0,
"num_beams": 1,
"num_return_sequences": 1,
"past_present_share_buffer": use_decoder_masked_mha,
"repetition_penalty": 1.0,
"temperature": 1.0,
"top_k": 1,
"top_p": 1.0,
},
}
# Requirements for the DMMHA kernel:
# - Buffer sharing = true
# - New input: past_sequence_length
# - New input: cache_indirection
# Otherwise, buffer sharing should be false and the new inputs should not be added
# for beam search to work in ORT GenAI.
if use_decoder_masked_mha:
genai_config["model"]["decoder"]["inputs"].update(
{
"past_sequence_length": "past_sequence_length",
"cache_indirection": "cache_indirection",
}
)
if output_qk:
genai_config["model"]["decoder"]["outputs"].update(
{
"output_cross_qk_names": "output_cross_qk_%d",
}
)
with open(os.path.join(output_dir, "genai_config.json"), "w") as f:
json.dump(genai_config, f, indent=4)
@staticmethod
def load_model(
model_name_or_path: str,
model_impl: str,
cache_dir: str,
device: torch.device,
dtype: torch.dtype,
merge_encoder_and_decoder_init: bool = True,
no_beam_search_op: bool = False,
output_qk: bool = False,
) -> dict[str, torch.nn.Module]:
"""Load model given a pretrained name or path, then build models for ONNX conversion.
Args:
model_name_or_path (str): pretrained model name or path
model_impl (str): library to load model from
cache_dir (str): cache directory
device (torch.device): device to run the model
dtype (torch.dtype): dtype to run the model
merge_encoder_and_decoder_init (bool, optional): Whether merge encoder and decoder initialization into one ONNX model. Defaults to True.
no_beam_search_op (bool, optional): Whether to use beam search op or not. Defaults to False.
output_qk (bool, optional): Whether to output QKs to calculate batched jump times for word-level timestamps. Defaults to False.
Returns:
Dict[str, torch.nn.Module]: mapping from name to modules for ONNX conversion.
"""
# Load PyTorch model
if model_impl == "hf":
# Load from Hugging Face
model = WhisperForConditionalGeneration.from_pretrained(
model_name_or_path, cache_dir=cache_dir, attn_implementation="eager"
)
else:
# Load from OpenAI
import whisper # noqa: PLC0415
if not os.path.exists(model_name_or_path):
name_or_path = model_name_or_path.split("/")[-1][8:]
else:
name_or_path = model_name_or_path
model = whisper.load_model(name_or_path, device, download_root=cache_dir, in_memory=True)
# Set PyTorch model properties
model.eval().to(device=device)
if model_impl == "hf":
model.to(dtype=dtype)
config = WhisperConfig.from_pretrained(model_name_or_path, cache_dir=cache_dir)
# Load each component of PyTorch model
decoder = WhisperDecoder(config, model, model_impl, no_beam_search_op).eval()
components = {"decoder": decoder}
if merge_encoder_and_decoder_init:
encoder_decoder_init = WhisperEncoderDecoderInit(config, model, model_impl, no_beam_search_op).eval()
components.update({"encoder": encoder_decoder_init})
else:
encoder = WhisperEncoder(config, model, model_impl).eval()
components.update({"encoder": encoder, "decoder_init": decoder})
if output_qk:
batched_jump_times = WhisperJumpTimes(config, device, cache_dir).eval()
components.update({"jump_times": batched_jump_times})
return components
@staticmethod
def export_onnx(
model: WhisperEncoder | WhisperEncoderDecoderInit | WhisperDecoder,
onnx_model_path: str,
provider: str,
verbose: bool,
use_external_data_format: bool,
use_fp16_inputs: bool,
use_int32_inputs: bool,
use_encoder_hidden_states: bool,
use_kv_cache_inputs: bool,
):
"""Export model component to ONNX
Args:
model (class): PyTorch class to export
onnx_model_path (str): path to save ONNX model
provider (str): provider to use for verifying parity on ONNX model
verbose (bool): print verbose information.
use_external_data_format (bool): use external data format or not.
use_fp16_inputs (bool): use float16 inputs for the audio_features, encoder_hidden_states, logits, and KV caches.
use_int32_inputs (bool): use int32 inputs for the decoder_input_ids.
use_encoder_hidden_states (bool): use encoder_hidden_states as model input for decoder-init/decoder-without-past models.
use_kv_cache_inputs (bool): use KV caches as model inputs for decoder-with-past models.
"""
if isinstance(model, WhisperEncoder):
model.export_onnx(
onnx_model_path,
provider,
verbose,
use_external_data_format,
use_fp16_inputs,
)
elif isinstance(model, WhisperEncoderDecoderInit):
model.export_onnx(
onnx_model_path,
provider,
verbose,
use_external_data_format,
use_fp16_inputs,
use_int32_inputs,
)
elif isinstance(model, WhisperDecoder):
model.export_onnx(
onnx_model_path,
provider,
verbose,
use_external_data_format,
use_fp16_inputs,
use_int32_inputs,
use_encoder_hidden_states,
use_kv_cache_inputs,
)
elif isinstance(model, WhisperJumpTimes):
model.export_onnx(
onnx_model_path,
provider,
verbose,
use_external_data_format,
use_fp16_inputs,
use_int32_inputs,
)
else:
raise ValueError(f"Unknown instance for model detected: {type(model)}")
@staticmethod
def optimize_onnx(
onnx_model_path: str,
optimized_model_path: str,
is_float16: bool,
num_attention_heads: int,
hidden_size: int,
num_decoder_layers: int,
use_external_data_format: bool = False,
use_gpu: bool = False,
provider: str = "cpu",
is_decoder: bool = False,
no_beam_search_op: bool = False,
use_decoder_masked_mha: bool = False,
output_qk: bool = False,
):
"""Optimize ONNX model with an option to convert it to use mixed precision."""
from fusion_options import FusionOptions # noqa: PLC0415
optimization_options = FusionOptions("bart")
optimization_options.use_multi_head_attention = True
optimization_options.disable_multi_head_attention_bias = provider == "rocm"
m = optimize_model(
onnx_model_path,
model_type="bart",
num_heads=num_attention_heads,
hidden_size=hidden_size,
opt_level=0,
optimization_options=optimization_options,
use_gpu=use_gpu,
only_onnxruntime=False,
)
# Add `past_sequence_length`, `cache_indirection`, and `output_qk` to `MultiHeadAttention` ops
if is_decoder and no_beam_search_op:
if use_decoder_masked_mha:
# FP16 CUDA, FP32 CUDA, and FP32 CPU use the `DecoderMaskedMultiHeadAttention` kernel
# via `MultiHeadAttention`, which requires the `past_sequence_length` and
# `cache_indirection` inputs
m, past_seq_len_name = fix_past_sequence_length(m)
m = add_cache_indirection_to_mha(m, past_seq_len_name)
if output_qk:
m = add_output_qk_to_mha(m, skip_node_idxs=list(range(0, 2 * num_decoder_layers, 2)))
m.save_model_to_file(optimized_model_path, use_external_data_format, all_tensors_to_one_file=True)
@staticmethod
def pt_transcription_for_verify_onnx(
processor: WhisperProcessor,
pt_model: torch.nn.Module,
device: torch.device,
batch_size: int = 1,
prompt_mode: bool = False,
):
# Try to import `datasets` pip package
try:
from datasets import load_dataset # noqa: PLC0415
except Exception as e:
logger.error(f"An error occurred while importing `datasets`: {e}", exc_info=True) # noqa: G201
install_cmd = "pip install datasets"
logger.warning(f"Could not import `datasets`. Attempting to install `datasets` via `{install_cmd}`.")
os.system(install_cmd)
from datasets import load_dataset # noqa: PLC0415
ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")
input_features_ = []
if batch_size == 1:
input_features = processor([ds[0]["audio"]["array"]], return_tensors="pt").input_features
else:
input_features_ = [
processor([ds[3]["audio"]["array"]], return_tensors="pt").input_features,
processor([ds[3]["audio"]["array"]], return_tensors="pt").input_features,
]
assert len(input_features_) == batch_size
input_features = torch.cat((input_features_[0], input_features_[1]))
max_length, min_length, num_beams, num_return_sequences = 30, 0, 1, 1
length_penalty, repetition_penalty = 1.0, 1.0
inputs = {
"input_features": input_features.to(device),
"max_length": max_length,
"min_length": min_length,
"num_beams": num_beams,
"num_return_sequences": num_return_sequences,
"length_penalty": length_penalty,
"repetition_penalty": repetition_penalty,
"early_stopping": True,
"use_cache": True,
}
if prompt_mode:
prompts = ["John has doubts", "Maria has grave doubts"]
prompt_ids = [processor.get_prompt_ids(p) for p in prompts]
pt_transcription = []
pt_outputs = []
# The looping for model.generate is necessary here due to the limitation as per
# https://huggingface.co/docs/transformers/model_doc/whisper#transformers.WhisperForConditionalGeneration.generate.prompt_ids
# prompt_ids input requires a tensor of rank 1
for i in range(batch_size):
inputs["prompt_ids"] = torch.from_numpy(prompt_ids[i]).to(device=device)
inputs["input_features"] = input_features_[i].to(device)
pt_output = pt_model.generate(**inputs).detach().cpu().numpy()
pt_outputs.append(pt_output)
pt_transcription.append(processor.batch_decode(pt_output, skip_special_tokens=True)[0])
inputs["input_features"] = input_features
del inputs["prompt_ids"]
else:
prompt_ids = []
pt_outputs = pt_model.generate(**inputs).detach().cpu().numpy()
pt_transcription = [processor.batch_decode(pt_outputs, skip_special_tokens=True)[0]]
pt_outputs = list(pt_outputs)
del inputs["early_stopping"]
del inputs["use_cache"]
return inputs, pt_transcription, pt_outputs, prompt_ids
@staticmethod
def select_transcription_options(
batch_size: int,
prompt_mode: bool,
):
if batch_size > 1 and prompt_mode:
expected_transcription_no_comma_prompt1 = " John has doubts whether Sir Frederick Layton's work is really Greek after all and can discover in it but little of Rocky I"
expected_transcription_misspelled_prompt1 = " John has doubts whether Sir Frederick Latins work is really Greek after all and can discover in it but little of Rocky I"
expected_transcription_no_comma_prompt2 = " Maria has grave doubts whether Sir Frederick Layton's work is really Greek after all and can discover in it but little of Rocky"
expected_transcription_misspelled_prompt2 = " Maria has grave doubts whether Sir Frederick Latins work is really Greek after all and can discover in it but little of Rocky I"
expected_transcription_options = {
expected_transcription_no_comma_prompt1,
expected_transcription_no_comma_prompt2,
expected_transcription_misspelled_prompt1,
expected_transcription_misspelled_prompt2,
}
else:
expected_transcription_no_comma = (
" Mr. Quilter is the apostle of the middle classes and we are glad to welcome his gospel."
)
expected_transcription_with_comma = (
" Mr. Quilter is the apostle of the middle classes, and we are glad to welcome his gospel."
)
expected_transcription_with_quote_and_comma = (
' "Mr. Quilter is the apostle of the middle classes, and we are glad to welcome his gospel.'
)
expected_transcription_options = {
expected_transcription_no_comma,
expected_transcription_with_comma,
expected_transcription_with_quote_and_comma,
}
return expected_transcription_options
@staticmethod
def get_outputs(
pt_outputs: np.ndarray,
ort_outputs: np.ndarray,
i: int,
):
"""Get PyTorch and ONNX Runtime output token ids at index i"""
pt_output, ort_output = pt_outputs[i], ort_outputs[i]
pt_shape, ort_shape = pt_output.shape, ort_output.shape
# Hugging Face impl. + Beam Search op: PyTorch = (26,) and ORT = (30,)
# OpenAI impl. + Beam Search op: PyTorch = (1, 30) and ORT = (30,)
if pt_shape != ort_shape:
if len(pt_shape) > 1:
pt_output = pt_output[0]
pt_shape = pt_output.shape
if len(ort_shape) > 1:
ort_output = ort_output[0]
ort_shape = ort_output.shape
if pt_shape[0] != ort_shape[0]:
min_len = min(pt_shape[0], ort_shape[0])
pt_output = pt_output[:min_len]
ort_output = ort_output[:min_len]
assert pt_output.shape == ort_output.shape
return pt_output, ort_output
@staticmethod
def verify_onnx(
model_name_or_path: str,
cache_dir: str,
ort_session: InferenceSession,
device: torch.device,
batch_size: int = 1,
prompt_mode: bool = False,
):
"""Compare the result from PyTorch and ONNX Runtime to verify the ONNX model is good."""
pt_model = WhisperForConditionalGeneration.from_pretrained(
model_name_or_path, cache_dir=cache_dir, attn_implementation="eager"
).to(device)
processor = WhisperProcessor.from_pretrained(model_name_or_path, cache_dir=cache_dir)
config = WhisperConfig.from_pretrained(model_name_or_path, cache_dir=cache_dir)
inputs, pt_transcription, pt_outputs, decoder_prompt_ids = WhisperHelper.pt_transcription_for_verify_onnx(
processor,
pt_model,
device,
batch_size=batch_size,
prompt_mode=prompt_mode,
)
start_id = [config.decoder_start_token_id] # ex: [50258]
prompt_ids = processor.get_decoder_prompt_ids(language="english", task="transcribe")
prompt_ids = [token[1] for token in prompt_ids] # ex: [50259, 50358, 50363]
forced_decoder_ids = start_id + prompt_ids # ex: [50258, 50259, 50358, 50363]
ort_names = [entry.name for entry in ort_session.get_inputs()]
ort_dtypes = [entry.type for entry in ort_session.get_inputs()]
ort_to_np = {
"tensor(float)": np.float32,
"tensor(float16)": np.float16,
"tensor(int64)": np.int64,
"tensor(int32)": np.int32,
"tensor(int8)": np.int8,
"tensor(uint8)": np.uint8,
}
use_extra_decoding_ids = "extra_decoding_ids" in ort_names
for name, dtype in zip(ort_names, ort_dtypes, strict=False):
if name == "input_features":
inputs[name] = inputs[name].detach().cpu().numpy()
elif name == "vocab_mask":
inputs[name] = np.ones(config.vocab_size, dtype=ort_to_np[dtype])
elif name == "prefix_vocab_mask":
inputs[name] = np.ones((batch_size, config.vocab_size), dtype=ort_to_np[dtype])
elif name == "decoder_input_ids":
if not prompt_mode:
raw_input_ids = [start_id] if use_extra_decoding_ids else [forced_decoder_ids]
inputs[name] = np.array(raw_input_ids, dtype=ort_to_np[dtype])
else:
# This logic handles the scenario for when prompts are not of the same size
# For example if our prompt ids are [p1_id_1, p1_id_2] and [p2_id_1]
# The final decoder_input_ids will look as such after padding
# [prev_token, p1_id_1, p1_id_2, start_token, lang_token, transcribe_token]
# [prev_token, p2_id_1, PAD_TOKEN, start_token, lang_token, transcribe_token]
ort_prompts = []
for i in range(batch_size):
ort_prompts.append(decoder_prompt_ids[i].tolist())
max_len = max(len(p) for p in ort_prompts)
padded_prompts = []
for p in ort_prompts:
padded_prompt = [*p, *([config.pad_token_id] * (max_len - len(p)))]
padded_prompts.append(padded_prompt + forced_decoder_ids)
inputs[name] = np.array(padded_prompts, dtype=ort_to_np[dtype])
elif name == "logits_processor":
inputs[name] = np.array([1], dtype=ort_to_np[dtype])
elif name == "cross_qk_layer_head":
inputs[name] = np.array([[0, 0]], dtype=ort_to_np[dtype])
elif name == "extra_decoding_ids":
inputs[name] = np.repeat(np.array([prompt_ids], dtype=ort_to_np[dtype]), batch_size, 0)
elif name == "temperature":
inputs[name] = np.array([1.0], dtype=ort_to_np[dtype])
else:
inputs[name] = np.array([inputs[name]], dtype=ort_to_np[dtype])
ort_outputs = ort_session.run(None, inputs)[0][:, 0, :]
ort_transcription = processor.batch_decode(ort_outputs, skip_special_tokens=True)
expected_transcription_options = WhisperHelper.select_transcription_options(batch_size, prompt_mode)
parity = 1
for i in range(batch_size):
pt_output, ort_output = WhisperHelper.get_outputs(pt_outputs, ort_outputs, i)
# Check if token ids match
parity *= np.allclose(pt_output, ort_output)
# Check if transcribed outputs match
parity *= (
pt_transcription[i] in expected_transcription_options
and ort_transcription[i] in expected_transcription_options
)
max_diff = 0
if not parity:
for i in range(batch_size):
pt_output, ort_output = WhisperHelper.get_outputs(pt_outputs, ort_outputs, i)
diff = pt_output - ort_output
max_diff_i = max(diff.min(), diff.max(), key=abs)
max_diff = max(max_diff, max_diff_i)
if max_diff != 0:
logger.warning(f"PyTorch outputs: {pt_transcription}")
logger.warning(f"ONNX Runtime outputs: {ort_transcription}")
return 0