| import os, json, torch |
| from torch.export import Dim |
| from transformers import MoonshineForConditionalGeneration, AutoProcessor, AutoConfig |
|
|
| MODEL = os.environ.get('MS_MODEL', 'moonshine-ai/moonshine-tiny-uk') |
| OUT = os.environ.get('MS_OUT', '/work/out') |
| m = MoonshineForConditionalGeneration.from_pretrained(MODEL, attn_implementation='eager').eval() |
|
|
| class Encoder(torch.nn.Module): |
| def __init__(self, m): super().__init__(); self.enc = m.model.encoder |
| def forward(self, input_values): return self.enc(input_values).last_hidden_state |
|
|
| audio = torch.zeros(1, 16000 * 3) |
| enc = Encoder(m).eval() |
| torch.onnx.export(enc, (audio,), f'{OUT}/encoder.onnx', dynamo=True, |
| input_names=['input_values'], output_names=['encoder_hidden_states'], |
| dynamic_shapes={'input_values': {1: Dim('samples', min=2000, max=16000*30)}}) |
| print('encoder exported') |
|
|
| proc = AutoProcessor.from_pretrained(MODEL) |
| vocab = proc.tokenizer.get_vocab() |
| with open(f'{OUT}/tokens.txt', 'w', encoding='utf-8') as f: |
| for piece, idx in sorted(vocab.items(), key=lambda kv: kv[1]): |
| f.write(f'{idx}\t{piece}\n') |
| cfg = AutoConfig.from_pretrained(MODEL) |
| json.dump({'decoder_start_token_id': cfg.decoder_start_token_id, 'eos_token_id': cfg.eos_token_id, |
| 'bos_token_id': cfg.bos_token_id, 'vocab_size': cfg.vocab_size, 'sampling_rate': 16000}, |
| open(f'{OUT}/decode_config.json', 'w'), indent=2) |
| print('tokens.txt:', len(vocab), 'entries') |
|
|
