# /// script
# requires-python = ">=3.12"
# dependencies = [
#     "onnx>=1.21.0",
#     "onnxruntime>=1.26.0",
#     "onnxscript>=0.7.0",
#     "torch>=2.12.0",
#     "torchvision>=0.27.0",
# ]
# ///
from pathlib import Path
import torch

# Load the uploaded PyTorch module
model_path = "embedding.pt"
import onnxruntime as ort


model = torch.load(model_path, map_location="cpu",weights_only=False)
model.eval()

# Create a dummy input for export
dummy_input = torch.randint(
    0,
    #config.num_audio_codebook * config.audio_vocab_size
     8 * 1025 ,
    (1, 8, 128),
    dtype=torch.long
)

# Export to ONNX
onnx_path = "embedding.onnx"
torch.onnx.export(
    model,
    dummy_input,
    onnx_path,export_params=True,
    #input_names=["input_ids"],
    # output_names=["embeddings"],
    # dynamic_axes={
    #     "input_ids": {0: "batch_size", 1: "sequence_length"},
    #     "embeddings": {0: "batch_size", 1: "sequence_length"},
    # },
    external_data=False,
    opset_version=19,
)

print(f"ONNX model saved to: {onnx_path}")
print(f"Embedding shape: ({model.num_embeddings}, {model.embedding_dim})")
session = ort.InferenceSession(onnx_path)

print(session.get_inputs())
print(session.get_outputs())


###
model_path = "linear_head.pt"
 


model = torch.load(model_path, map_location="cpu",weights_only=False)
model.eval()

# Create a dummy input for export
dummy_input = torch.randn(1, 1024, dtype=torch.float32)

# Export to ONNX
onnx_path = "liner_head.onnx"
torch.onnx.export(
    model,
    dummy_input,
    onnx_path,
    export_params=True,
    input_names=["in_features"],
    output_names=["out_features"],
    external_data=False,
    # dynamic_axes={
    #     "input_ids": {0: "batch_size", 1: "sequence_length"},
    #     "embeddings": {0: "batch_size", 1: "sequence_length"},
    # },
    opset_version=19,
)

print(f"ONNX model saved to: {onnx_path}")
print(f"Embedding shape: ({model.in_features}, {model.out_features})")
session = ort.InferenceSession(onnx_path)

print(session.get_inputs())
print(session.get_outputs())