Skip to content

Commit 3a0e935

Browse files
zhuzilinzhuzilin
committed
finish exporting .en models
1 parent df8cf85 commit 3a0e935

File tree

2 files changed

+58
-31
lines changed

2 files changed

+58
-31
lines changed

‎whisper/decoding.py

Lines changed: 27 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -133,41 +133,56 @@ def __init__(self, model: "Whisper", initial_token_length: int):
133133
self.model: "Whisper" = model
134134
self.initial_token_length = initial_token_length
135135
self.kv_cache = None
136-
self.export_onnx = False
136+
if model.type == "tiny.en":
137+
self.kv_cache_size = lambda x, y: [8, x, y, 384]
138+
elif model.type == "base.en":
139+
self.kv_cache_size = lambda x, y: [12, x, y, 512]
140+
elif model.type == "small.en":
141+
self.kv_cache_size = lambda x, y: [24, x, y, 768]
142+
elif model.type == "medium.en":
143+
self.kv_cache_size = lambda x, y: [48, x, y, 1024]
144+
else:
145+
raise ValueError(f"Unsupported model type: {model.type}")
137146

138147
def logits(self, tokens: Tensor, audio_features: Tensor) -> Tensor:
148+
n_group = tokens.shape[0]
139149
if self.kv_cache is None:
140-
# hard code for decoder layer 4, 6, 8, 10
141-
self.kv_cache = np.zeros([8, 5, self.initial_token_length, 384], dtype=np.float32)
150+
self.kv_cache = np.zeros(
151+
self.kv_cache_size(n_group, self.initial_token_length), dtype=np.float32)
142152
offset = 0
143153
else:
144154
offset = self.kv_cache.shape[2]
145-
new_kv_cache = np.zeros([8, 5, offset + 1, 384], dtype=np.float32)
155+
new_kv_cache = np.zeros(self.kv_cache_size(n_group, offset + 1), dtype=np.float32)
146156
new_kv_cache[:, :, :-1, :] = self.kv_cache
147157
self.kv_cache = new_kv_cache
148158

149159
if tokens.shape[-1] > self.initial_token_length:
150160
# only need to use the last token except in the first forward pass
151161
tokens = tokens[:, -1:]
152162

153-
if self.export_onnx and self.kv_cache.shape[2] > self.initial_token_length:
163+
# export decoder as onnx
164+
if False and self.kv_cache.shape[2] > self.initial_token_length:
165+
print(f"tokens: {tokens.shape}")
166+
print(f"audio_features: {audio_features.shape}")
167+
print(f"kv_cache: {self.kv_cache.shape}")
154168
torch.onnx.export(
155169
self.model.decoder,
156170
(tokens, audio_features, torch.from_numpy(self.kv_cache), torch.tensor(offset)),
157171
"decoder.onnx",
158-
verbose=True,
172+
verbose=False,
159173
opset_version=13,
160174
input_names=["tokens", "audio_features", "kv_cache", "offset"],
161175
output_names=["logits", "output_kv_cache"],
162176
dynamic_axes={
163-
"tokens": [1],
164-
"kv_cache": [2],
177+
"tokens": [0, 1],
178+
"audio_features": [0],
179+
"kv_cache": [1, 2],
165180
"output_kv_cache": [2],
166181
}
167182
)
168183
exit()
169-
output, self.kv_cache = self.model.decoder(tokens, audio_features, kv_cache=self.kv_cache, offset=offset)
170-
# output, self.kv_cache = self.model.decoder(tokens, audio_features, kv_cache=torch.from_numpy(self.kv_cache), offset=torch.tensor(offset))
184+
#output, self.kv_cache = self.model.decoder(tokens, audio_features, kv_cache=self.kv_cache, offset=offset)
185+
output, self.kv_cache = self.model.decoder(tokens, audio_features, kv_cache=torch.from_numpy(self.kv_cache), offset=torch.tensor(offset))
171186
return output
172187

173188
def cleanup_caching(self):
@@ -578,6 +593,7 @@ def _get_audio_features(self, mel: Tensor):
578593
# encoded audio features are given; skip audio encoding
579594
audio_features = mel
580595
else:
596+
# # export encoder as onnx
581597
# torch.onnx.export(
582598
# self.model.encoder,
583599
# (mel),
@@ -615,6 +631,7 @@ def _main_loop(self, audio_features: Tensor, tokens: Tensor):
615631
try:
616632
for i in range(self.sample_len):
617633
logits = self.inference.logits(tokens, audio_features)
634+
print(f"step: {i}, logits: {logits}", flush=True)
618635

619636
if i == 0 and self.tokenizer.no_speech is not None: # save no_speech_probs
620637
probs_at_sot = logits[:, self.sot_index].float().softmax(dim=-1)

‎whisper/model.py

Lines changed: 31 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -84,7 +84,12 @@ def forward(
8484
k = self.key(x if xa is None else xa)
8585
v = self.value(x if xa is None else xa)
8686
if kv_cache is not None and k.shape[1] <= self.n_ctx:
87-
key_id = self.layer_id - 4
87+
# here is hard coded
88+
# tiny.en: 4
89+
# base.en: 6
90+
# small.en: 12
91+
# medium.en: 24
92+
key_id = self.layer_id - 24
8893
value_id = key_id + 1
8994
size = k.shape[1]
9095
kv_cache[key_id, :, -size:, :] = k
@@ -210,8 +215,10 @@ def __init__(self, model: str):
210215

211216
self.core = Core()
212217
self._model = self.core.read_model(
213-
hf_hub_download(repo_id=f"zhuzilin/whisper-openvino-{model}", filename="encoder.xml"),
214-
hf_hub_download(repo_id=f"zhuzilin/whisper-openvino-{model}", filename="encoder.bin"),
218+
# hf_hub_download(repo_id=f"zhuzilin/whisper-openvino-{model}", filename="encoder.xml"),
219+
# hf_hub_download(repo_id=f"zhuzilin/whisper-openvino-{model}", filename="encoder.bin"),
220+
"encoder.xml",
221+
"encoder.bin",
215222
)
216223
self.model = self.core.compile_model(self._model, "CPU")
217224

@@ -226,8 +233,10 @@ def __init__(self, model: str):
226233

227234
self.core = Core()
228235
self._model = self.core.read_model(
229-
hf_hub_download(repo_id=f"zhuzilin/whisper-openvino-{model}", filename="decoder.xml"),
230-
hf_hub_download(repo_id=f"zhuzilin/whisper-openvino-{model}", filename="decoder.bin"),
236+
# hf_hub_download(repo_id=f"zhuzilin/whisper-openvino-{model}", filename="decoder.xml"),
237+
# hf_hub_download(repo_id=f"zhuzilin/whisper-openvino-{model}", filename="decoder.bin"),
238+
"decoder.xml",
239+
"decoder.bin",
231240
)
232241
self.model = self.core.compile_model(self._model, "CPU")
233242

@@ -246,23 +255,24 @@ def forward(self, x: Tensor, xa: Tensor, kv_cache: Tensor, offset: int):
246255
class Whisper(nn.Module):
247256
def __init__(self, dims: ModelDimensions, model: str):
248257
super().__init__()
258+
self.type = model
249259
self.dims = dims
250-
# self.encoder = AudioEncoder(
251-
# self.dims.n_mels,
252-
# self.dims.n_audio_ctx,
253-
# self.dims.n_audio_state,
254-
# self.dims.n_audio_head,
255-
# self.dims.n_audio_layer,
256-
# )
257-
# self.decoder = TextDecoder(
258-
# self.dims.n_vocab,
259-
# self.dims.n_text_ctx,
260-
# self.dims.n_text_state,
261-
# self.dims.n_text_head,
262-
# self.dims.n_text_layer,
263-
# )
264-
self.encoder = OpenVinoAudioEncoder(model=model)
265-
self.decoder = OpenVinoTextDecoder(model=model)
260+
self.encoder = AudioEncoder(
261+
self.dims.n_mels,
262+
self.dims.n_audio_ctx,
263+
self.dims.n_audio_state,
264+
self.dims.n_audio_head,
265+
self.dims.n_audio_layer,
266+
)
267+
self.decoder = TextDecoder(
268+
self.dims.n_vocab,
269+
self.dims.n_text_ctx,
270+
self.dims.n_text_state,
271+
self.dims.n_text_head,
272+
self.dims.n_text_layer,
273+
)
274+
# self.encoder = OpenVinoAudioEncoder(model=model)
275+
# self.decoder = OpenVinoTextDecoder(model=model)
266276

267277
def embed_audio(self, mel: torch.Tensor):
268278
return self.encoder.forward(mel)

0 commit comments

Comments
 (0)