import torch
import onnx
import os
from train import PaletteGenerator

def export_to_onnx():
    print("--- 1. Loading Model ---")
    device = torch.device("cpu")
    model = PaletteGenerator().to(device)
    model.load_state_dict(torch.load("model.pth", map_location=device))
    model.eval()

    dummy_input_ids = torch.randint(0, 1000, (1, 10), dtype=torch.long).to(device)
    dummy_mask = torch.ones((1, 10), dtype=torch.long).to(device) 

    temp_filename = "palette_model_temp.onnx"
    final_filename = "palette_model.onnx"
    
    print(f"\n--- 2. Exporting ---")
    torch.onnx.export(
        model, 
        (dummy_input_ids, dummy_mask), 
        temp_filename,
        export_params=True,
        opset_version=17, 
        do_constant_folding=True,
        input_names=['input_ids', 'attention_mask'],
        output_names=['output'],
        dynamic_axes={
            'input_ids': {0: 'batch_size', 1: 'sequence_length'},
            'attention_mask': {0: 'batch_size', 1: 'sequence_length'},
            'output': {0: 'batch_size'}
        }
    )
    
    print(f"\n--- 3. Merging External Weights ---")
    model_proto = onnx.load(temp_filename)
    onnx.save(model_proto, final_filename)
    
    size_mb = os.path.getsize(final_filename) / (1024 * 1024)
    print(f"Final File: {final_filename}")
    print(f"Size: {size_mb:.2f} MB")
    
if __name__ == "__main__":
    export_to_onnx()