"""Compare encoder.onnx against the torch encoder+projector for several audio lengths."""

import numpy as np
import onnxruntime as rt
import soundfile as sf
import torch
from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor

OUT = "/media/hdd16/onnx-asr-exports/granite-speech-3.3-2b"
torch.set_num_threads(8)
torch.set_grad_enabled(False)

processor = AutoProcessor.from_pretrained("ibm-granite/granite-speech-3.3-2b")
model = AutoModelForSpeechSeq2Seq.from_pretrained(
    "ibm-granite/granite-speech-3.3-2b", dtype=torch.float32, device_map="cpu"
).eval()

opts = rt.SessionOptions()
opts.intra_op_num_threads = 8
sess = rt.InferenceSession(f"{OUT}/encoder.onnx", opts, providers=["CPUExecutionProvider"])

wav, _ = sf.read(f"{OUT}/clips/en_2.wav", dtype="float32")
for n in (16000, 16000 + 137, 48000, len(wav)):
    x = wav[:n][None]
    feats = processor.audio_processor(torch.from_numpy(x), device="cpu")
    ref = model.projector(model.encoder(feats["input_features"])).numpy()
    got = sess.run(["audio_embeds"], {"input_features": x})[0]
    exp_len = int(feats["input_features_mask"].sum())
    print(
        f"n={n:7d} onnx={got.shape} torch={ref.shape} expected_len={exp_len} "
        f"max_abs={np.abs(got - ref).max():.3e}"
    )