import torch
from torch.export import Dim
from transformers import AutoProcessor, MoonshineForConditionalGeneration

MODEL, OUT = 'moonshine-ai/moonshine-tiny-uk', '/work/out'
model = MoonshineForConditionalGeneration.from_pretrained(MODEL, attn_implementation='eager').eval()
cfg = model.config
print('params:', round(sum(p.numel() for p in model.parameters())/1e6, 1), 'M',
      '| start:', cfg.decoder_start_token_id, '| eos:', cfg.eos_token_id)

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

class Decoder(torch.nn.Module):
    def __init__(self, m):
        super().__init__(); self.m = m
    def forward(self, input_ids, encoder_hidden_states):
        out = self.m.model.decoder(input_ids=input_ids,
                                   encoder_hidden_states=encoder_hidden_states,
                                   use_cache=False)
        return self.m.proj_out(out.last_hidden_state)

audio = torch.zeros(1, 16000 * 3)
enc = Encoder(model).eval()
with torch.no_grad():
    hs = enc(audio)
print('encoder out:', tuple(hs.shape))

samples = Dim('samples', min=2000, max=16000 * 30)
frames = Dim('frames', min=4, max=2000)
tokens = Dim('tokens', min=1, max=512)

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: samples}})
print('encoder exported')

dec = Decoder(model).eval()
ids = torch.tensor([[cfg.decoder_start_token_id, 100, 200]], dtype=torch.long)
with torch.no_grad():
    logits = dec(ids, hs)
print('decoder out:', tuple(logits.shape))

torch.onnx.export(dec, (ids, hs), f'{OUT}/decoder.onnx', dynamo=True,
                  input_names=['input_ids', 'encoder_hidden_states'], output_names=['logits'],
                  dynamic_shapes={'input_ids': {1: tokens},
                                  'encoder_hidden_states': {1: frames}})
print('decoder exported')