#!/usr/bin/env python3 """Export CohereLabs/cohere-transcribe-arabic-07-2026 to the onnx-community merged layout that WinSTT's native Rust CohereEngine consumes: onnx/encoder_model.onnx input_features[B,T,128] -> last_hidden_state[B,S,1024] onnx/decoder_model_merged.onnx input_ids, attention_mask, position_ids, num_logits_to_keep, encoder_hidden_states, past_key_values.{i}.{decoder,encoder}.{key,value} -> logits, present.{i}.{decoder,encoder}.{key,value} (i=0..7) Module wiring adapted from eschmidbauer/cohere-transcribe-03-2026-onnx/export_onnx.py (proven on the same architecture), restructured so the graph I/O names/shapes match onnx-community/cohere-transcribe- 03-2026-ONNX byte-for-byte (verified against that repo's declared inputs/outputs). """ import os from pathlib import Path import onnx import torch import torch.nn as nn import torch.nn.functional as F from transformers import AutoModelForSpeechSeq2Seq MODEL_DIR = "model_ar" OUT = Path("onnx") OPSET = 21 NUM_DECODER_LAYERS = 8 # --------------------------------------------------------------------------- encoder class EncoderForExport(nn.Module): """Full conformer encoder as ONE graph: (B,T,128) time-first mel -> (B,S,1024).""" def __init__(self, model): super().__init__() self.pre_encode = model.encoder.pre_encode self.pos_enc = model.encoder.pos_enc self.layers = model.encoder.layers self.encoder_decoder_proj = model.encoder_decoder_proj self.pos_enc._materialize_pe( length=self.pos_enc.max_len, device=torch.device("cpu"), dtype=torch.float32 ) def _create_masks(self, length, max_len): att_mask = torch.ones(1, max_len, max_len, dtype=torch.bool, device=length.device) pad_mask = torch.arange(0, max_len, device=length.device).expand( length.size(0), -1 ) < length.unsqueeze(-1) pad_att = pad_mask.unsqueeze(1).repeat([1, max_len, 1]) pad_att = torch.logical_and(pad_att, pad_att.transpose(1, 2)) att_mask = torch.logical_and(att_mask, pad_att) return ~pad_mask, ~att_mask def forward(self, input_features): # onnx-community feeds (B, T, 128) time-first; pre_encode wants (B, 128, T) channel-first. x_in = input_features.transpose(1, 2) length = torch.full( (x_in.size(0),), x_in.size(2), dtype=torch.long, device=x_in.device ) x, length = self.pre_encode(x_in, length) length = length.to(torch.int64) max_len = x.size(1) x, pos_emb = self.pos_enc(x) pad_mask, att_mask = self._create_masks(length, max_len) for layer in self.layers: x = layer(x, pos_emb, mask=att_mask, pad_mask=pad_mask) if self.encoder_decoder_proj is not None: x = self.encoder_decoder_proj(x) return x # last_hidden_state (B, S, 1024) # --------------------------------------------------------------------------- decoder class DecoderMergedForExport(nn.Module): """Merged autoregressive decoder. Cross-attn K/V are recomputed from encoder_hidden_states every step and echoed as present.*.encoder (full length); the Rust loop stores them once and the recompute keeps them correct regardless of the (declared-but-unused) past.*.encoder inputs.""" def __init__(self, model): super().__init__() self.embedding = model.transf_decoder._embedding self.decoder_layers = model.transf_decoder._decoder.layers self.final_layer_norm = model.transf_decoder._decoder.final_layer_norm self.classifier = model.log_softmax def forward(self, input_ids, attention_mask, position_ids, num_logits_to_keep, encoder_hidden_states, *past): # past is flat: [dec_k0, dec_v0, enc_k0, enc_v0, dec_k1, ...]. The encoder (cross) past entries # are recomputed from encoder_hidden_states below, so they are functionally unused — but the Rust # engine binds all 32 past inputs and index-zips past<->present, so they MUST stay in the graph # signature. `keep_enc` (a scalar 0) references them so torch.onnx does not prune them. past_dec_k = [past[4 * i + 0] for i in range(NUM_DECODER_LAYERS)] past_dec_v = [past[4 * i + 1] for i in range(NUM_DECODER_LAYERS)] keep_enc = sum(past[4 * i + 2].sum() + past[4 * i + 3].sum() for i in range(NUM_DECODER_LAYERS)) * 0.0 B, tgt_len = input_ids.shape dtype = encoder_hidden_states.dtype past_len = past_dec_k[0].shape[2] total_kv = past_len + tgt_len q_pos = torch.arange(past_len, past_len + tgt_len, device=input_ids.device)[:, None] k_pos = torch.arange(total_kv, device=input_ids.device)[None, :] self_mask = torch.zeros((B, 1, tgt_len, total_kv), device=input_ids.device, dtype=dtype) self_mask.masked_fill_((k_pos > q_pos)[None, None], float("-inf")) # padding mask from attention_mask (B, total_kv): additive -inf on padded keys (no-op when all 1s). pad = (1 - attention_mask)[:, None, None, :].to(dtype) * torch.finfo(dtype).min self_mask = self_mask + pad hidden = self.embedding(input_ids, position_ids) present = [] # flat: dec_k, dec_v, enc_k, enc_v per layer for i, layer in enumerate(self.decoder_layers): attn_self = layer.first_sub_layer attn_cross = layer.second_sub_layer # self-attention with growing KV cache residual = hidden h = layer.layer_norm_1(hidden) q = attn_self._reshape(attn_self.query_net(h)) k = torch.cat([past_dec_k[i], attn_self._reshape(attn_self.key_net(h))], dim=2) v = torch.cat([past_dec_v[i], attn_self._reshape(attn_self.value_net(h))], dim=2) out = F.scaled_dot_product_attention(q, k, v, attn_mask=self_mask, scale=attn_self.scale) out = out.transpose(1, 2).contiguous().view(B, tgt_len, attn_self.hidden_size) hidden = residual + attn_self.out_projection(out) # cross-attention: recompute K/V from encoder states (echoed as present.encoder) ck = attn_cross.key_net(encoder_hidden_states).view( B, -1, attn_cross.num_heads, attn_cross.head_dim).transpose(1, 2) cv = attn_cross.value_net(encoder_hidden_states).view( B, -1, attn_cross.num_heads, attn_cross.head_dim).transpose(1, 2) residual = hidden h = layer.layer_norm_2(hidden) q = attn_cross._reshape(attn_cross.query_net(h)) out = F.scaled_dot_product_attention(q, ck, cv, attn_mask=None, scale=attn_cross.scale) out = out.transpose(1, 2).contiguous().view(B, tgt_len, attn_cross.hidden_size) hidden = residual + attn_cross.out_projection(out) # FFN hidden = hidden + layer.third_sub_layer(layer.layer_norm_3(hidden)) present.extend([k, v, ck, cv]) logits = self.classifier(self.final_layer_norm(hidden)) + keep_enc # (B, T, 16384) start = logits.shape[1] - num_logits_to_keep # keep last num_logits_to_keep rows logits = logits[:, start:, :] return (logits, *present) # --------------------------------------------------------------------------- export helpers def _save(model, path): data_name = Path(path).name + "_data" onnx.save(model, path, save_as_external_data=True, all_tensors_to_one_file=True, location=data_name, convert_attribute=True) def export_encoder(model): wr = EncoderForExport(model).eval() dummy = torch.randn(1, 500, 128) path = str(OUT / "encoder_model.onnx") torch.onnx.export( wr, (dummy,), path, opset_version=OPSET, dynamo=False, input_names=["input_features"], output_names=["last_hidden_state"], dynamic_axes={"input_features": {0: "batch_size", 1: "sequence_length"}, "last_hidden_state": {0: "batch_size", 1: "encoder_sequence_length"}}, ) m = onnx.load(path, load_external_data=True) _save(m, path) print(" encoder ok") def export_decoder(model): wr = DecoderMergedForExport(model).eval() H, D, src, past_len = 8, 128, 62, 10 input_ids = torch.ones(1, 1, dtype=torch.long) attention_mask = torch.ones(1, past_len + 1, dtype=torch.long) position_ids = torch.tensor([[past_len]], dtype=torch.long) num_logits_to_keep = torch.tensor(1, dtype=torch.long) enc = torch.randn(1, src, 1024) past = [] for _ in range(NUM_DECODER_LAYERS): past += [torch.randn(1, H, past_len, D), torch.randn(1, H, past_len, D), # decoder k,v torch.randn(1, H, src, D), torch.randn(1, H, src, D)] # encoder k,v (unused) args = (input_ids, attention_mask, position_ids, num_logits_to_keep, enc, *past) input_names = ["input_ids", "attention_mask", "position_ids", "num_logits_to_keep", "encoder_hidden_states"] output_names = ["logits"] dyn = { "input_ids": {0: "batch_size", 1: "sequence_length"}, "attention_mask": {0: "batch_size", 1: "total_sequence_length"}, "position_ids": {0: "batch_size", 1: "sequence_length"}, "encoder_hidden_states": {0: "batch_size", 1: "encoder_sequence_length"}, "logits": {0: "batch_size", 1: "num_logits_to_keep"}, } for i in range(NUM_DECODER_LAYERS): for kind, ax in (("decoder", "past_decoder_sequence_length"), ("encoder", "past_encoder_sequence_length")): for kv in ("key", "value"): n = f"past_key_values.{i}.{kind}.{kv}" input_names.append(n) dyn[n] = {0: "batch_size", 2: ax} for i in range(NUM_DECODER_LAYERS): for kind, ax in (("decoder", "total_decoder_sequence_length"), ("encoder", "present_encoder_sequence_length")): for kv in ("key", "value"): n = f"present.{i}.{kind}.{kv}" output_names.append(n) dyn[n] = {0: "batch_size", 2: ax} path = str(OUT / "decoder_model_merged.onnx") torch.onnx.export( wr, args, path, opset_version=OPSET, dynamo=False, input_names=input_names, output_names=output_names, dynamic_axes=dyn, ) m = onnx.load(path, load_external_data=True) _save(m, path) print(" decoder ok") def main(): OUT.mkdir(exist_ok=True) print(f"Loading {MODEL_DIR}") model = AutoModelForSpeechSeq2Seq.from_pretrained( MODEL_DIR, trust_remote_code=True, dtype=torch.float32 ).eval() with torch.no_grad(): print("Exporting encoder..."); export_encoder(model) print("Exporting decoder..."); export_decoder(model) print("Done ->", OUT.resolve()) if __name__ == "__main__": main()