import torch
from onnxruntime.transformers import optimizer
from onnxruntime.transformers.fusion_options import FusionOptions
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer

model_name = "hf-internal-testing/tiny-random-bart"
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

model.eval()


class EncoderWrapper(torch.nn.Module):
    def __init__(self, encoder: torch.nn.Module):
        super().__init__()
        self.encoder = encoder

    def forward(
        self, input_ids: torch.Tensor, attention_mask: torch.Tensor
    ) -> torch.Tensor:
        outs = self.encoder(input_ids, attention_mask)
        return outs["last_hidden_state"]


model = EncoderWrapper(encoder=model.model.encoder)
print(model)

text = "Hello, how are you?"
inputs = tokenizer(text, return_tensors="pt")

input_ids = inputs["input_ids"]
attention_mask = inputs["attention_mask"]

input_names = ["input_ids", "attention_mask"]
output_names = ["encoder_output"]

onnx_path = "bart_model.onnx"

torch.onnx.export(
    model,
    (input_ids, attention_mask),
    onnx_path,
    export_params=True,
    input_names=input_names,
    output_names=output_names,
    dynamic_axes={
        "input_ids": {0: "batch_size", 1: "sequence_length"},
        "attention_mask": {0: "batch_size", 1: "sequence_length"},
        "encoder_output": {0: "batch_size", 1: "sequence_length"},
    },
    opset_version=16,
)
print(f"BART encoder exported to {onnx_path}")

optimization_options = FusionOptions("bart")
optimization_options.enable_attention = True


m = optimizer.optimize_model(
    onnx_path,
    model_type="bart",
    num_heads=0,
    hidden_size=0,
    opt_level=2,
    use_gpu=False,
    verbose=True,
    optimization_options=optimization_options,
    only_onnxruntime=False,
)

optimized_path = "bart_encoder_optimized.onnx"
m.save_model_to_file(optimized_path)

print(f"Optimized ONNX model saved to {optimized_path}")
print(m.get_fused_operator_statistics())

import torch
from onnxruntime.transformers import optimizer
from transformers import AutoTokenizer, BertModel

model = BertModel.from_pretrained("hf-internal-testing/tiny-random-bert")
tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-bert")

model.eval()

text = "Hello, how are you?"
inputs = tokenizer(text, return_tensors="pt")
input_names = ["input_ids", "attention_mask"]
output_names = ["output"]

onnx_path = "bert_model.onnx"

torch.onnx.export(
    model,
    (inputs["input_ids"], inputs["attention_mask"]),
    onnx_path,
    input_names=input_names,
    output_names=output_names,
    dynamic_axes={
        "input_ids": {0: "batch_size", 1: "sequence_length"},
        "attention_mask": {0: "batch_size", 1: "sequence_length"},
        "output": {0: "batch_size", 1: "sequence_length"},
    },
    opset_version=20,
    verbose=True,
)
print(f"Model exported to {onnx_path}")

m = optimizer.optimize_model(
    onnx_path,
    model_type="bert",
    num_heads=0,
    hidden_size=0,
    opt_level=2,
    use_gpu="cpu",
    verbose=True,
)

m.save_model_to_file(
    "bert_model_optimized.onnx",
)
print(m.get_fused_operator_statistics())