""" Export Cohere Transcribe to ONNX (encoder + decoder) matching sherpa-onnx Whisper tensor contract. Requirements: pip install torch transformers safetensors sentencepiece onnx onnxruntime librosa soundfile protobuf Usage: python export_onnx.py --model-dir . --output-dir ./onnx-export Memory: ~8GB RAM peak (loads model in bfloat16, traces on CPU). Works on 24GB shared RAM with no dedicated GPU. """ import argparse import json import math import os from pathlib import Path import numpy as np import torch import torch.nn as nn # --------------------------------------------------------------------------- # Encoder wrapper: mel features -> pre-computed cross K/V per decoder layer # # This reshapes Cohere's architecture to match sherpa-onnx's Whisper contract: # Inputs: mel (batch, n_mels, T) # Outputs: n_layer_cross_k (n_dec_layers, batch, T', n_state) # n_layer_cross_v (n_dec_layers, batch, T', n_state) # --------------------------------------------------------------------------- class EncoderForONNX(nn.Module): """Wraps Cohere's ConformerEncoder + projection + cross-attn K/V pre-computation.""" def __init__(self, model): super().__init__() self.encoder = model.encoder self.encoder_decoder_proj = model.encoder_decoder_proj # 1280 -> 1024, may be None # Grab the cross-attention key/value projections from each decoder layer. # In the original model, these run inside the decoder on every step. # We pre-compute them once in the encoder to match Whisper's contract. dec_layers = model.transf_decoder._decoder.layers self.cross_k_projs = nn.ModuleList([layer.second_sub_layer.key_net for layer in dec_layers]) self.cross_v_projs = nn.ModuleList([layer.second_sub_layer.value_net for layer in dec_layers]) self.num_heads = dec_layers[0].second_sub_layer.num_heads self.head_dim = dec_layers[0].second_sub_layer.head_dim def forward(self, mel): # mel: (batch, 128, T) -- log-mel features, already preprocessed batch = mel.shape[0] # Encoder expects (batch, feat_in, T) and an optional length tensor. # For ONNX export we pass full-length (no padding). length = torch.full((batch,), mel.shape[2], dtype=torch.long, device=mel.device) enc_out, _enc_len = self.encoder(mel, length) # enc_out: (batch, T', 1280) if self.encoder_decoder_proj is not None: enc_out = self.encoder_decoder_proj(enc_out) # enc_out: (batch, T', 1024) # Pre-compute cross-attention K/V for each decoder layer # Sherpa-onnx Whisper expects: (n_layers, batch, T', n_state) cross_k_list = [] cross_v_list = [] for k_proj, v_proj in zip(self.cross_k_projs, self.cross_v_projs): cross_k_list.append(k_proj(enc_out)) cross_v_list.append(v_proj(enc_out)) # Stack: (n_layers, batch, T', hidden_size) n_layer_cross_k = torch.stack(cross_k_list, dim=0) n_layer_cross_v = torch.stack(cross_v_list, dim=0) return n_layer_cross_k, n_layer_cross_v # --------------------------------------------------------------------------- # Decoder wrapper: tokens + caches -> logits + updated caches # # Matches sherpa-onnx Whisper decoder contract: # Inputs: tokens (batch, n_tokens) # in_n_layer_self_k_cache (n_layers, batch, ctx_len, n_state) # in_n_layer_self_v_cache (n_layers, batch, ctx_len, n_state) # n_layer_cross_k (n_layers, batch, T', n_state) # n_layer_cross_v (n_layers, batch, T', n_state) # offset (scalar int64) -- positional offset # Outputs: logits (batch, n_tokens, vocab_size) # out_n_layer_self_k_cache (n_layers, batch, ctx_len, n_state) # out_n_layer_self_v_cache (n_layers, batch, ctx_len, n_state) # --------------------------------------------------------------------------- class DecoderLayerForONNX(nn.Module): """A single decoder layer that takes explicit K/V instead of running projections.""" def __init__(self, original_layer, layer_idx): super().__init__() self.layer_idx = layer_idx # Self-attention components self.layer_norm_1 = original_layer.layer_norm_1 self.self_attn_q = original_layer.first_sub_layer.query_net self.self_attn_k = original_layer.first_sub_layer.key_net self.self_attn_v = original_layer.first_sub_layer.value_net self.self_attn_out = original_layer.first_sub_layer.out_projection self.num_heads = original_layer.first_sub_layer.num_heads self.head_dim = original_layer.first_sub_layer.head_dim self.scale = original_layer.first_sub_layer.scale # Cross-attention: only need query projection + output projection. # K/V are pre-computed in encoder. self.layer_norm_2 = original_layer.layer_norm_2 self.cross_attn_q = original_layer.second_sub_layer.query_net self.cross_attn_out = original_layer.second_sub_layer.out_projection # Feed-forward self.layer_norm_3 = original_layer.layer_norm_3 self.ffn = original_layer.third_sub_layer def _reshape(self, x): b, t, _ = x.shape return x.view(b, t, self.num_heads, self.head_dim).transpose(1, 2) def forward( self, hidden_states, # (batch, n_tokens, hidden_size) cross_k, # (batch, T', hidden_size) -- pre-computed for this layer cross_v, # (batch, T', hidden_size) -- pre-computed for this layer in_self_k_cache, # (batch, n_heads, ctx_len, head_dim) in_self_v_cache, # (batch, n_heads, ctx_len, head_dim) self_attention_mask, # (batch, 1, n_tokens, total_kv_len) offset, # scalar int64 ): # --- Self-attention with KV cache --- residual = hidden_states hidden_states = self.layer_norm_1(hidden_states) q = self._reshape(self.self_attn_q(hidden_states)) k = self._reshape(self.self_attn_k(hidden_states)) v = self._reshape(self.self_attn_v(hidden_states)) # Update self-attention cache: write new K/V at the offset position n_tokens = k.shape[2] out_self_k_cache = in_self_k_cache.clone() out_self_v_cache = in_self_v_cache.clone() out_self_k_cache[:, :, offset:offset + n_tokens, :] = k out_self_v_cache[:, :, offset:offset + n_tokens, :] = v # Use the full cache up to offset + n_tokens as keys/values total_kv_len = offset + n_tokens cached_k = out_self_k_cache[:, :, :total_kv_len, :] cached_v = out_self_v_cache[:, :, :total_kv_len, :] attn_out = torch.nn.functional.scaled_dot_product_attention( q, cached_k, cached_v, attn_mask=self_attention_mask[:, :, :, :total_kv_len], dropout_p=0.0, scale=self.scale, ) attn_out = attn_out.transpose(1, 2).contiguous().view( hidden_states.shape[0], hidden_states.shape[1], self.num_heads * self.head_dim ) hidden_states = residual + self.self_attn_out(attn_out) # --- Cross-attention (K/V pre-computed) --- residual = hidden_states hidden_states = self.layer_norm_2(hidden_states) cross_q = self._reshape(self.cross_attn_q(hidden_states)) # cross_k and cross_v need reshaping to (batch, heads, T', head_dim) b = cross_k.shape[0] cross_k_reshaped = cross_k.view(b, -1, self.num_heads, self.head_dim).transpose(1, 2) cross_v_reshaped = cross_v.view(b, -1, self.num_heads, self.head_dim).transpose(1, 2) cross_attn_out = torch.nn.functional.scaled_dot_product_attention( cross_q, cross_k_reshaped, cross_v_reshaped, dropout_p=0.0, scale=self.scale, ) cross_attn_out = cross_attn_out.transpose(1, 2).contiguous().view( hidden_states.shape[0], hidden_states.shape[1], self.num_heads * self.head_dim ) hidden_states = residual + self.cross_attn_out(cross_attn_out) # --- Feed-forward --- residual = hidden_states hidden_states = self.layer_norm_3(hidden_states) hidden_states = residual + self.ffn(hidden_states) return hidden_states, out_self_k_cache, out_self_v_cache class DecoderForONNX(nn.Module): """Wraps Cohere's TransformerDecoder to match sherpa-onnx Whisper decoder contract.""" def __init__(self, model): super().__init__() dec_wrapper = model.transf_decoder # Embedding + positional encoding self.token_embedding = dec_wrapper._embedding.token_embedding self.position_embedding = dec_wrapper._embedding.position_embedding self.embedding_layer_norm = dec_wrapper._embedding.layer_norm # Decoder layers (rewrapped for explicit K/V) self.layers = nn.ModuleList([ DecoderLayerForONNX(layer, i) for i, layer in enumerate(dec_wrapper._decoder.layers) ]) self.final_layer_norm = dec_wrapper._decoder.final_layer_norm # Classification head (weights tied to token_embedding) self.classifier = model.log_softmax self.num_layers = len(self.layers) self.hidden_size = dec_wrapper._embedding.token_embedding.embedding_dim def forward( self, tokens, # (batch, n_tokens) int64 in_n_layer_self_k_cache, # (n_layers, batch, n_heads, ctx_len, head_dim) in_n_layer_self_v_cache, # (n_layers, batch, n_heads, ctx_len, head_dim) n_layer_cross_k, # (n_layers, batch, T', hidden_size) n_layer_cross_v, # (n_layers, batch, T', hidden_size) offset, # scalar int64 ): batch, n_tokens = tokens.shape # Positional IDs from offset position_ids = torch.arange(offset, offset + n_tokens, device=tokens.device) position_ids = position_ids.unsqueeze(0).expand(batch, -1) # Embeddings hidden_states = self.embedding_layer_norm( self.token_embedding(tokens) + self.position_embedding(position_ids) ) # Build causal self-attention mask total_kv_len = offset + n_tokens query_positions = torch.arange(offset, offset + n_tokens, device=tokens.device).unsqueeze(1) key_positions = torch.arange(total_kv_len, device=tokens.device).unsqueeze(0) causal_mask = key_positions > query_positions # True = masked self_attention_mask = torch.zeros( (batch, 1, n_tokens, total_kv_len), device=tokens.device, dtype=hidden_states.dtype ) self_attention_mask.masked_fill_(causal_mask.unsqueeze(0).unsqueeze(0), float("-inf")) # Run through decoder layers out_k_caches = [] out_v_caches = [] for i, layer in enumerate(self.layers): hidden_states, out_k, out_v = layer( hidden_states=hidden_states, cross_k=n_layer_cross_k[i], cross_v=n_layer_cross_v[i], in_self_k_cache=in_n_layer_self_k_cache[i], in_self_v_cache=in_n_layer_self_v_cache[i], self_attention_mask=self_attention_mask, offset=offset, ) out_k_caches.append(out_k) out_v_caches.append(out_v) hidden_states = self.final_layer_norm(hidden_states) logits = self.classifier(hidden_states) out_n_layer_self_k_cache = torch.stack(out_k_caches, dim=0) out_n_layer_self_v_cache = torch.stack(out_v_caches, dim=0) return logits, out_n_layer_self_k_cache, out_n_layer_self_v_cache # --------------------------------------------------------------------------- # tokens.txt generation # --------------------------------------------------------------------------- def generate_tokens_txt(model_dir, output_path): """Generate sherpa-onnx tokens.txt from the SentencePiece tokenizer.""" import sentencepiece as spm tokenizer_config_path = os.path.join(model_dir, "tokenizer_config.json") with open(tokenizer_config_path, "r", encoding="utf-8") as f: tokenizer_config = json.load(f) sp = spm.SentencePieceProcessor() sp.Load(os.path.join(model_dir, "tokenizer.model")) # Start with SentencePiece vocab vocab = {} for i in range(sp.get_piece_size()): token = sp.id_to_piece(i) vocab[i] = token # Override with added_tokens_decoder (these have the correct special token mappings) added_tokens = tokenizer_config.get("added_tokens_decoder", {}) for idx_str, info in added_tokens.items(): idx = int(idx_str) vocab[idx] = info["content"] # Write in sherpa-onnx format: "token id" per line max_id = max(vocab.keys()) with open(output_path, "w", encoding="utf-8") as f: for i in range(max_id + 1): token = vocab.get(i, f"") f.write(f"{token} {i}\n") print(f"Wrote {max_id + 1} tokens to {output_path}") # --------------------------------------------------------------------------- # Main export logic # --------------------------------------------------------------------------- def export_encoder(model, output_dir, opset_version=17): """Export encoder to ONNX.""" print("Wrapping encoder...") encoder = EncoderForONNX(model) encoder.eval() # Dummy input: (batch=1, n_mels=128, T=1000) ~= 10 seconds of audio dummy_mel = torch.randn(1, 128, 1000, dtype=torch.float32) output_path = os.path.join(output_dir, "cohere-encoder.onnx") print(f"Exporting encoder to {output_path}...") with torch.no_grad(): torch.onnx.export( encoder, (dummy_mel,), output_path, input_names=["mel"], output_names=["n_layer_cross_k", "n_layer_cross_v"], dynamic_axes={ "mel": {0: "batch", 2: "T"}, "n_layer_cross_k": {1: "batch", 2: "T_enc"}, "n_layer_cross_v": {1: "batch", 2: "T_enc"}, }, opset_version=opset_version, do_constant_folding=True, dynamo=False, ) print(f"Encoder exported: {output_path}") return output_path def export_decoder(model, output_dir, opset_version=17): """Export decoder to ONNX.""" print("Wrapping decoder...") decoder = DecoderForONNX(model) decoder.eval() config = model.config n_dec_layers = config.transf_decoder["config_dict"]["num_layers"] # 8 hidden_size = config.transf_decoder["config_dict"]["hidden_size"] # 1024 n_heads = config.transf_decoder["config_dict"]["num_attention_heads"] # 8 head_dim = hidden_size // n_heads # 128 max_ctx = config.transf_decoder["config_dict"]["max_sequence_length"] # 1024 # Dummy inputs batch = 1 n_tokens = 4 # prompt length for first step T_enc = 125 # encoder output length for ~10s audio offset = torch.tensor(0, dtype=torch.int64) dummy_tokens = torch.zeros(batch, n_tokens, dtype=torch.int64) dummy_self_k_cache = torch.zeros(n_dec_layers, batch, n_heads, max_ctx, head_dim) dummy_self_v_cache = torch.zeros(n_dec_layers, batch, n_heads, max_ctx, head_dim) dummy_cross_k = torch.zeros(n_dec_layers, batch, T_enc, hidden_size) dummy_cross_v = torch.zeros(n_dec_layers, batch, T_enc, hidden_size) output_path = os.path.join(output_dir, "cohere-decoder.onnx") print(f"Exporting decoder to {output_path}...") with torch.no_grad(): torch.onnx.export( decoder, (dummy_tokens, dummy_self_k_cache, dummy_self_v_cache, dummy_cross_k, dummy_cross_v, offset), output_path, input_names=[ "tokens", "in_n_layer_self_k_cache", "in_n_layer_self_v_cache", "n_layer_cross_k", "n_layer_cross_v", "offset", ], output_names=[ "logits", "out_n_layer_self_k_cache", "out_n_layer_self_v_cache", ], dynamic_axes={ "tokens": {0: "batch", 1: "n_tokens"}, "in_n_layer_self_k_cache": {1: "batch"}, "in_n_layer_self_v_cache": {1: "batch"}, "n_layer_cross_k": {1: "batch", 2: "T_enc"}, "n_layer_cross_v": {1: "batch", 2: "T_enc"}, "logits": {0: "batch", 1: "n_tokens"}, "out_n_layer_self_k_cache": {1: "batch"}, "out_n_layer_self_v_cache": {1: "batch"}, }, opset_version=opset_version, do_constant_folding=True, dynamo=False, ) print(f"Decoder exported: {output_path}") return output_path def main(): parser = argparse.ArgumentParser(description="Export Cohere Transcribe to ONNX for sherpa-onnx") parser.add_argument("--model-dir", type=str, default=".", help="Directory with model files") parser.add_argument("--output-dir", type=str, default="./onnx-export", help="Output directory for ONNX files") parser.add_argument("--opset", type=int, default=17, help="ONNX opset version") parser.add_argument("--skip-encoder", action="store_true", help="Skip encoder export") parser.add_argument("--skip-decoder", action="store_true", help="Skip decoder export") parser.add_argument("--skip-tokens", action="store_true", help="Skip tokens.txt generation") args = parser.parse_args() os.makedirs(args.output_dir, exist_ok=True) # Generate tokens.txt first (no model loading needed) if not args.skip_tokens: print("=== Generating tokens.txt ===") generate_tokens_txt(args.model_dir, os.path.join(args.output_dir, "tokens.txt")) if args.skip_encoder and args.skip_decoder: print("Both encoder and decoder skipped. Done.") return # Load model in bfloat16 to save memory, then convert to float32 for ONNX export print("=== Loading model ===") print("Loading in bfloat16 to save memory...") from transformers import AutoModelForSpeechSeq2Seq model = AutoModelForSpeechSeq2Seq.from_pretrained( args.model_dir, trust_remote_code=True, dtype=torch.bfloat16, ) model.eval() # Convert to float32 for ONNX compatibility (ONNX doesn't support bfloat16 well) print("Converting to float32 for ONNX export...") model = model.float() if not args.skip_encoder: print("\n=== Exporting Encoder ===") export_encoder(model, args.output_dir, args.opset) if not args.skip_decoder: print("\n=== Exporting Decoder ===") export_decoder(model, args.output_dir, args.opset) print("\n=== Done ===") print(f"Output files in: {args.output_dir}/") print("Next steps:") print(" 1. Validate with: python validate_onnx.py --input-dir ./onnx-export") print(" 2. Quantize with: python quantize.py --input-dir ./onnx-export") print(" 3. Use in C# with sherpa-onnx (see CohereTranscribe.cs)") if __name__ == "__main__": main()