| """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}" |
| ) |
|
|
