#!/usr/bin/env python3
"""Local smoke test for the distilled HeAR ViT-S Canon model package."""

from __future__ import annotations

import argparse
from pathlib import Path

import torch
from transformers import AutoModel


def main() -> None:
  ap = argparse.ArgumentParser(description="Smoke test for local HF upload directory.")
  ap.add_argument("--model-dir", type=Path, default=Path(__file__).resolve().parent)
  ap.add_argument("--batch-size", type=int, default=4)
  args = ap.parse_args()

  model = AutoModel.from_pretrained(str(args.model_dir), trust_remote_code=True)
  model.eval()

  raw_audio = torch.rand((int(args.batch_size), 32000), dtype=torch.float32)
  with torch.inference_mode():
    out_from_wave = model(input_values=raw_audio, return_dict=True).pooler_output

  spectrogram = model.preprocess_audio(raw_audio)
  with torch.inference_mode():
    out_from_spec = model(pixel_values=spectrogram, return_dict=True).pooler_output

  max_abs = (out_from_wave - out_from_spec).abs().max().item()

  print(f"model_dir={args.model_dir}")
  print(f"spectrogram_shape={tuple(spectrogram.shape)}")
  print(f"wave_embedding_shape={tuple(out_from_wave.shape)}")
  print(f"spec_embedding_shape={tuple(out_from_spec.shape)}")
  print(f"max_abs_diff_wave_vs_spec={max_abs:.8f}")


if __name__ == "__main__":
  main()