Views
No views yet
huggingface-hub Python library.pip3 install huggingface-hub>=0.17.1huggingface-cli download Esperanto/whisper-large-v3-kvc-fp16-onnx --local-dir whisper-large-v3-kvc-fp16-onnx --local-dir-use-symlinks Falsehuggingface-cli, please see: HF -> Hub Python Library -> Download files -> Download from the CLI.1#!/usr/bin/env python3
2import whisper
3import onnx
4import sys
5import time
6import onnxruntime
7from typing import Sequence, Optional
8import numpy as np
9from pathlib import Path
10
11def run_whisper_decoder(decoder_model_path, execution_provider, session_options, decoder_output_names, cross_attn_tensors, num_new_tokens, provider_options = {}):
12 start = time.time()
13 decoder_session = onnxruntime.InferenceSession(decoder_model_path, sess_options=session_options, providers=[execution_provider], provider_options=[provider_options])
14 compile_time = time.time()
15 transcription = decoder_loop(decoder_session, decoder_output_names, cross_attn_tensors, num_new_tokens)
16 inference_time = time.time()
17 return transcription
18
19
20def decoder_loop(decoder_session, decoder_output_names, cross_attn_tensors, num_new_tokens):
21 # Generate start of transcription tokens
22 tokenizer = whisper.tokenizer.get_tokenizer(multilingual=True)
23 first_tokens = np.array([tokenizer.sot, 0, tokenizer.transcribe, tokenizer.no_timestamps], dtype=np.int64)
24
25 # Self attention mask key, value vectors
26 self_attn_past_k = []
27 self_attn_past_v = []
28 for i in range(32):
29 self_attn_past_k.append(np.zeros((1, 20, 447, 64), dtype=np.float16))
30 self_attn_past_v.append(np.zeros((1, 20, 447, 64), dtype=np.float16))
31
32 # Cross attention
33 cross_attn_k = cross_attn_tensors[0::2]
34 cross_attn_v = cross_attn_tensors[1::2]
35
36 # Attention mask
37 attn_mask_size = 448
38 attn_mask = np.zeros((1,attn_mask_size), dtype=np.int64)
39
40 # Process first tokens
41 for j in range(len(first_tokens)):
42 tokens = np.array([first_tokens[j]], dtype=np.int64).reshape(1, 1)
43 attn_mask[0,-1 - j] = 1
44
45 decoder_input = {"input_ids": tokens, "attention_mask": attn_mask}
46 for i in range(32):
47 decoder_input[f"past_key_values.{str(i)}.key"] = self_attn_past_k[i]
48 decoder_input[f"past_key_values.{str(i)}.value"] = self_attn_past_v[i]
49 decoder_input[f"cross_attn.{str(i)}.key"] = cross_attn_k[i]
50 decoder_input[f"cross_attn.{str(i)}.value"] = cross_attn_v[i]
51
52 logits, *cache_tensors = decoder_session.run(decoder_output_names, decoder_input)
53 next_token = np.argmax(logits[0,0])
54
55 self_attn_k = cache_tensors[0::2]
56 self_attn_v = cache_tensors[1::2]
57 for i in range(32):
58 self_attn_past_k[i] = self_attn_k[i][:,:,1:,:]
59 self_attn_past_v[i] = self_attn_v[i][:,:,1:,:]
60
61 if (j == 0):
62 # set language token
63 first_tokens[1] = next_token
64
65 transcribed_tokens = [next_token]
66 for j in range(4, 4 + num_new_tokens):
67 tokens = np.array([transcribed_tokens[-1]], dtype=np.int64).reshape(1, 1)
68 attn_mask[0,-1 - j] = 1
69
70 decoder_input = {"input_ids": tokens, "attention_mask": attn_mask}
71 for i in range(32):
72 decoder_input[f"past_key_values.{str(i)}.key"] = self_attn_past_k[i]
73 decoder_input[f"past_key_values.{str(i)}.value"] = self_attn_past_v[i]
74 decoder_input[f"cross_attn.{str(i)}.key"] = cross_attn_k[i]
75 decoder_input[f"cross_attn.{str(i)}.value"] = cross_attn_v[i]
76
77 logits, *cache_tensors = decoder_session.run(decoder_output_names, decoder_input)
78 next_token = np.argmax(logits[0,0])
79 # print(j, next_token)
80 if next_token == tokenizer.eot: # end_of_transcription
81 break
82 transcribed_tokens.append(next_token)
83 self_attn_k = cache_tensors[0::2]
84 self_attn_v = cache_tensors[1::2]
85 for i in range(32):
86 self_attn_past_k[i] = self_attn_k[i][:,:,1:,:]
87 self_attn_past_v[i] = self_attn_v[i][:,:,1:,:]
88
89 return tokenizer.decode(transcribed_tokens)
90
91
92def main(argv: Optional[Sequence[str]] = None):
93 num_seconds = 28.8
94
95 speech_path = 'sample_audio.wav'
96 encoder_model_path = 'whisper-large-v3-kvc-fp16-onnx/encoder/model.onnx'
97 decoder_model_path = 'whisper-large-v3-kvc-fp16-onnx/decoder/model.onnx'
98
99 # Load audio
100 print(f"Spectrogram speech audio file {speech_path}... ", end="")
101 audio = whisper.load_audio(speech_path)
102 audio = whisper.pad_or_trim(audio, length=int(num_seconds*16000))
103 mel = whisper.log_mel_spectrogram(audio, n_mels=128).unsqueeze(0) # Unsqueeze to set batch=1
104 print("OK")
105
106 print("Running encoder... ", end="")
107
108 # Session options
109 session_options = onnxruntime.SessionOptions()
110 # Disable all the graph optimizations
111 session_options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL
112
113 # Encode
114 encoder = onnx.load(encoder_model_path, load_external_data=False)
115 encoder_input = {"mel": mel.numpy().astype('float16')}
116 encoder_output_names = [tensor.name for tensor in encoder.graph.output]
117 # CPU encoding
118 cpu_provider = 'CPUExecutionProvider'
119 enc_session_cpu = onnxruntime.InferenceSession(encoder_model_path, sess_options=session_options, providers=[cpu_provider])
120 cross_attn_tensors_cpu = enc_session_cpu.run(encoder_output_names, encoder_input)
121
122 print("OK")
123
124 # DECODE API PARAMS
125 max_context = 448
126 new_tokens = 20
127
128 # Run decoder model CPU
129 decoder = onnx.load(decoder_model_path, load_external_data=False)
130 decoder_output_names = [tensor.name for tensor in decoder.graph.output]
131
132 run_whisper_decoder(decoder_model_path, cpu_provider, session_options, decoder_output_names, cross_attn_tensors_cpu, new_tokens)
133
134
135if __name__ == "__main__":
136 sys.exit(main(sys.argv[1:]))