import json, os, time, numpy as np, onnx, onnxruntime as ort
from onnxruntime.quantization import quantize_dynamic, QuantType
from transformers import AutoProcessor, AutoConfig

OUT = '/work/out'
for name in ('encoder', 'decoder'):
    src, dst = f'{OUT}/{name}.onnx', f'{OUT}/{name}.int8.onnx'
    quantize_dynamic(src, dst, weight_type=QuantType.QInt8, extra_options={'MatMulConstBOnly': True})
    print(f'{name}: fp32 {os.path.getsize(src) + os.path.getsize(src + ".data" if os.path.exists(src + ".data") else src):,}'
          f' -> int8 {os.path.getsize(dst):,} bytes')

# tokens.txt for a Dart-side decoder: id -> piece, one per line, tab separated
proc = AutoProcessor.from_pretrained('moonshine-ai/moonshine-tiny-uk')
tok = proc.tokenizer
vocab = tok.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')
print('tokens.txt:', len(vocab), 'entries,', os.path.getsize(f'{OUT}/tokens.txt'), 'bytes')
cfg = AutoConfig.from_pretrained('moonshine-ai/moonshine-tiny-uk')
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('decode_config.json written')