Initial upload of distilled ViT-S Canon model

README.md

@@ -0,0 +1,145 @@
+---
+license: other
+license_name: health-ai-developer-foundations
+license_link: https://developers.google.com/health-ai-developer-foundations/terms
+library_name: transformers
+pipeline_tag: feature-extraction
+tags:
+- audio
+- medical
+- embeddings
+- vision-transformer
+- distillation
+- canon
+---
+
+# Distilled HeAR ViT-S Canon (PyTorch)
+
+This repository package contains a Hugging Face-compatible export of our distilled HeAR student:
+
+- Backbone: ViT-S (`vit_small_patch16_224`), 1-channel input
+- Input shape: `[B, 1, 192, 128]` mel+PCEN spectrograms from 2 s audio at 16 kHz
+- Canon configuration: A/B/C/D enabled, 2D Canon (`kernel=4`), no positional encodings
+- Output embedding: `pooler_output` with shape `[B, 512]`
+
+This folder is ready for upload to Hugging Face Hub as-is.
+
+## Files in this package
+
+- `config.json`: model config + `auto_map` for custom loading
+- `configuration_hear_canon.py`: custom `PretrainedConfig`
+- `modeling_hear_canon.py`: custom `PreTrainedModel` with built-in audio preprocessing
+- `pytorch_model.bin`: distilled student + projection head weights
+- `preprocessor_config.json`: preprocessing metadata
+- `model_shapes.json`: parameter and tensor shape inventory
+- `training_args.json`: saved training/checkpoint args used for this export
+- `.gitattributes`: LFS patterns for Hub upload
+- `smoke_test.py`: local verification script
+
+## Quick start (local folder)
+
+```bash
+pip install -U "transformers>=4.50.0" timm torch scipy soundfile
+python3 trained_model_hf_upload/smoke_test.py
+```
+
+## Inference from raw audio waveform
+
+```python
+import torch
+from transformers import AutoModel
+
+model = AutoModel.from_pretrained(
+    "trained_model_hf_upload",
+    trust_remote_code=True,
+)
+model.eval()
+
+# 4 clips, each 2 seconds at 16 kHz => 32000 samples
+raw_audio_batch = torch.rand((4, 32000), dtype=torch.float32)
+
+with torch.inference_mode():
+    out = model(input_values=raw_audio_batch, return_dict=True)
+
+embeddings = out.pooler_output
+print(embeddings.shape)  # torch.Size([4, 512])
+```
+
+## Inference from `.wav` file
+
+```python
+import torch
+import soundfile as sf
+from scipy import signal
+from transformers import AutoModel
+
+
+def load_wav_mono_16k(path: str, target_sr: int = 16000) -> torch.Tensor:
+    audio, sr = sf.read(path, dtype="float32", always_2d=False)
+    if audio.ndim == 2:
+        audio = audio.mean(axis=1)
+    if sr != target_sr:
+        new_len = int(round(audio.shape[0] * (target_sr / sr)))
+        audio = signal.resample(audio, new_len)
+    return torch.from_numpy(audio).float()
+
+
+model = AutoModel.from_pretrained("trained_model_hf_upload", trust_remote_code=True)
+model.eval()
+
+waveform = load_wav_mono_16k("example.wav")
+
+with torch.inference_mode():
+    embedding = model.embed_audio(waveform)
+
+print(embedding.shape)  # torch.Size([1, 512])
+```
+
+## Inference from preprocessed spectrograms
+
+```python
+import torch
+from transformers import AutoModel
+
+model = AutoModel.from_pretrained("trained_model_hf_upload", trust_remote_code=True)
+model.eval()
+
+raw_audio = torch.rand((2, 32000), dtype=torch.float32)
+spectrogram = model.preprocess_audio(raw_audio)
+
+with torch.inference_mode():
+    out = model(pixel_values=spectrogram, return_dict=True)
+
+print(spectrogram.shape)         # torch.Size([2, 1, 192, 128])
+print(out.pooler_output.shape)   # torch.Size([2, 512])
+```
+
+## After uploading to Hugging Face
+
+Replace the local path with your Hub repo id:
+
+```python
+from transformers import AutoModel
+
+model = AutoModel.from_pretrained(
+    "<your-org>/<your-repo>",
+    trust_remote_code=True,
+)
+```
+
+Upload example:
+
+```bash
+huggingface-cli repo create <your-repo> --type model
+huggingface-cli upload <your-org>/<your-repo> trained_model_hf_upload .
+```
+
+## Architecture summary
+
+- Student parameters: `22,140,288`
+- Projection head parameters: `197,120`
+- Total parameters: `22,337,408`
+- Student hidden size: `384`
+- Pooler output size: `512`
+
+Detailed tensor shapes are listed in `model_shapes.json`.

config.json

@@ -0,0 +1,42 @@
+{
+  "_name_or_path": "distilled-hear-vit-s-canon",
+  "architectures": [
+    "HearCanonViTModel"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_hear_canon.HearCanonViTConfig",
+    "AutoModel": "modeling_hear_canon.HearCanonViTModel"
+  },
+  "clip_seconds": 2.0,
+  "canon": true,
+  "canon_2d": true,
+  "canon_a": true,
+  "canon_abcd": true,
+  "canon_b": true,
+  "canon_b_qkv": false,
+  "canon_c": true,
+  "canon_causal": false,
+  "canon_d": true,
+  "canon_kernel": 4,
+  "canon_no_pos_enc": true,
+  "hidden_act": "gelu",
+  "hidden_size": 384,
+  "image_size": [
+    192,
+    128
+  ],
+  "intermediate_size": 1536,
+  "layer_norm_eps": 1e-06,
+  "model_type": "hear_canon_vit",
+  "num_attention_heads": 6,
+  "num_audio_samples": 32000,
+  "num_channels": 1,
+  "num_hidden_layers": 12,
+  "patch_size": 16,
+  "pooled_dim": 512,
+  "pooler_output_size": 512,
+  "sample_rate": 16000,
+  "timm_model_name": "vit_small_patch16_224",
+  "torch_dtype": "float32",
+  "transformers_version": "4.50.3"
+}

configuration_hear_canon.py

@@ -0,0 +1,74 @@
+"""Configuration for the distilled HeAR ViT-S + Canon model."""
+
+from __future__ import annotations
+
+from typing import List
+
+from transformers import PretrainedConfig
+
+
+class HearCanonViTConfig(PretrainedConfig):
+  """Config class for the distilled HeAR ViT-S model with Canon layers."""
+
+  model_type = "hear_canon_vit"
+
+  def __init__(
+    self,
+    timm_model_name: str = "vit_small_patch16_224",
+    image_size: List[int] | tuple[int, int] = (192, 128),
+    patch_size: int = 16,
+    num_channels: int = 1,
+    hidden_size: int = 384,
+    num_hidden_layers: int = 12,
+    num_attention_heads: int = 6,
+    intermediate_size: int = 1536,
+    pooled_dim: int = 512,
+    pooler_output_size: int = 512,
+    hidden_act: str = "gelu",
+    layer_norm_eps: float = 1e-6,
+    sample_rate: int = 16000,
+    clip_seconds: float = 2.0,
+    num_audio_samples: int = 32000,
+    canon: bool = True,
+    canon_2d: bool = True,
+    canon_kernel: int = 4,
+    canon_a: bool = True,
+    canon_b: bool = True,
+    canon_b_qkv: bool = False,
+    canon_c: bool = True,
+    canon_d: bool = True,
+    canon_abcd: bool = True,
+    canon_no_pos_enc: bool = True,
+    canon_causal: bool = False,
+    **kwargs,
+  ):
+    self.timm_model_name = str(timm_model_name)
+    self.image_size = [int(image_size[0]), int(image_size[1])]
+    self.patch_size = int(patch_size)
+    self.num_channels = int(num_channels)
+    self.hidden_size = int(hidden_size)
+    self.num_hidden_layers = int(num_hidden_layers)
+    self.num_attention_heads = int(num_attention_heads)
+    self.intermediate_size = int(intermediate_size)
+    self.pooled_dim = int(pooled_dim)
+    self.pooler_output_size = int(pooler_output_size)
+    self.hidden_act = str(hidden_act)
+    self.layer_norm_eps = float(layer_norm_eps)
+
+    self.sample_rate = int(sample_rate)
+    self.clip_seconds = float(clip_seconds)
+    self.num_audio_samples = int(num_audio_samples)
+
+    self.canon = bool(canon)
+    self.canon_2d = bool(canon_2d)
+    self.canon_kernel = int(canon_kernel)
+    self.canon_a = bool(canon_a)
+    self.canon_b = bool(canon_b)
+    self.canon_b_qkv = bool(canon_b_qkv)
+    self.canon_c = bool(canon_c)
+    self.canon_d = bool(canon_d)
+    self.canon_abcd = bool(canon_abcd)
+    self.canon_no_pos_enc = bool(canon_no_pos_enc)
+    self.canon_causal = bool(canon_causal)
+
+    super().__init__(**kwargs)

model_shapes.json

@@ -0,0 +1,1353 @@
+{
+  "checkpoint_source": "/home/matt/Documents/HeAR/trained_model.pt",
+  "exported_weights_file": "pytorch_model.bin",
+  "input_spectrogram_shape": [
+    "B",
+    1,
+    192,
+    128
+  ],
+  "input_waveform_shape": [
+    "B",
+    32000
+  ],
+  "pooler_output_size": 512,
+  "projection_parameters": 197120,
+  "projection_state_dict_shapes": {
+    "bias": [
+      512
+    ],
+    "weight": [
+      512,
+      384
+    ]
+  },
+  "projection_state_dict_size": 2,
+  "student_hidden_size": 384,
+  "student_parameters": 22140288,
+  "student_state_dict_shapes": {
+    "blocks.0.block.attn.canon._fallback.conv.bias": [
+      384
+    ],
+    "blocks.0.block.attn.canon._fallback.conv.weight": [
+      384,
+      1,
+      4
+    ],
+    "blocks.0.block.attn.canon.conv.bias": [
+      384
+    ],
+    "blocks.0.block.attn.canon.conv.weight": [
+      384,
+      1,
+      4,
+      4
+    ],
+    "blocks.0.block.attn.module.proj.0.bias": [
+      384
+    ],
+    "blocks.0.block.attn.module.proj.0.weight": [
+      384,
+      384
+    ],
+    "blocks.0.block.attn.module.proj.1._fallback.conv.bias": [
+      384
+    ],
+    "blocks.0.block.attn.module.proj.1._fallback.conv.weight": [
+      384,
+      1,
+      4
+    ],
+    "blocks.0.block.attn.module.proj.1.conv.bias": [
+      384
+    ],
+    "blocks.0.block.attn.module.proj.1.conv.weight": [
+      384,
+      1,
+      4,
+      4
+    ],
+    "blocks.0.block.attn.module.qkv.bias": [
+      1152
+    ],
+    "blocks.0.block.attn.module.qkv.weight": [
+      1152,
+      384
+    ],
+    "blocks.0.block.mlp.canon._fallback.conv.bias": [
+      384
+    ],
+    "blocks.0.block.mlp.canon._fallback.conv.weight": [
+      384,
+      1,
+      4
+    ],
+    "blocks.0.block.mlp.canon.conv.bias": [
+      384
+    ],
+    "blocks.0.block.mlp.canon.conv.weight": [
+      384,
+      1,
+      4,
+      4
+    ],
+    "blocks.0.block.mlp.module.fc1.canon._fallback.conv.bias": [
+      1536
+    ],
+    "blocks.0.block.mlp.module.fc1.canon._fallback.conv.weight": [
+      1536,
+      1,
+      4
+    ],
+    "blocks.0.block.mlp.module.fc1.canon.conv.bias": [
+      1536
+    ],
+    "blocks.0.block.mlp.module.fc1.canon.conv.weight": [
+      1536,
+      1,
+      4,
+      4
+    ],
+    "blocks.0.block.mlp.module.fc1.fc1.bias": [
+      1536
+    ],
+    "blocks.0.block.mlp.module.fc1.fc1.weight": [
+      1536,
+      384
+    ],
+    "blocks.0.block.mlp.module.fc2.bias": [
+      384
+    ],
+    "blocks.0.block.mlp.module.fc2.weight": [
+      384,
+      1536
+    ],
+    "blocks.0.block.norm1.bias": [
+      384
+    ],
+    "blocks.0.block.norm1.weight": [
+      384
+    ],
+    "blocks.0.block.norm2.bias": [
+      384
+    ],
+    "blocks.0.block.norm2.weight": [
+      384
+    ],
+    "blocks.1.block.attn.canon._fallback.conv.bias": [
+      384
+    ],
+    "blocks.1.block.attn.canon._fallback.conv.weight": [
+      384,
+      1,
+      4
+    ],
+    "blocks.1.block.attn.canon.conv.bias": [
+      384
+    ],
+    "blocks.1.block.attn.canon.conv.weight": [
+      384,
+      1,
+      4,
+      4
+    ],
+    "blocks.1.block.attn.module.proj.0.bias": [
+      384
+    ],
+    "blocks.1.block.attn.module.proj.0.weight": [
+      384,
+      384
+    ],
+    "blocks.1.block.attn.module.proj.1._fallback.conv.bias": [
+      384
+    ],
+    "blocks.1.block.attn.module.proj.1._fallback.conv.weight": [
+      384,
+      1,
+      4
+    ],
+    "blocks.1.block.attn.module.proj.1.conv.bias": [
+      384
+    ],
+    "blocks.1.block.attn.module.proj.1.conv.weight": [
+      384,
+      1,
+      4,
+      4
+    ],
+    "blocks.1.block.attn.module.qkv.bias": [
+      1152
+    ],
+    "blocks.1.block.attn.module.qkv.weight": [
+      1152,
+      384
+    ],
+    "blocks.1.block.mlp.canon._fallback.conv.bias": [
+      384
+    ],
+    "blocks.1.block.mlp.canon._fallback.conv.weight": [
+      384,
+      1,
+      4
+    ],
+    "blocks.1.block.mlp.canon.conv.bias": [
+      384
+    ],
+    "blocks.1.block.mlp.canon.conv.weight": [
+      384,
+      1,
+      4,
+      4
+    ],
+    "blocks.1.block.mlp.module.fc1.canon._fallback.conv.bias": [
+      1536
+    ],
+    "blocks.1.block.mlp.module.fc1.canon._fallback.conv.weight": [
+      1536,
+      1,
+      4
+    ],
+    "blocks.1.block.mlp.module.fc1.canon.conv.bias": [
+      1536
+    ],
+    "blocks.1.block.mlp.module.fc1.canon.conv.weight": [
+      1536,
+      1,
+      4,
+      4
+    ],
+    "blocks.1.block.mlp.module.fc1.fc1.bias": [
+      1536
+    ],
+    "blocks.1.block.mlp.module.fc1.fc1.weight": [
+      1536,
+      384
+    ],
+    "blocks.1.block.mlp.module.fc2.bias": [
+      384
+    ],
+    "blocks.1.block.mlp.module.fc2.weight": [
+      384,
+      1536
+    ],
+    "blocks.1.block.norm1.bias": [
+      384
+    ],
+    "blocks.1.block.norm1.weight": [
+      384
+    ],
+    "blocks.1.block.norm2.bias": [
+      384
+    ],
+    "blocks.1.block.norm2.weight": [
+      384
+    ],
+    "blocks.10.block.attn.canon._fallback.conv.bias": [
+      384
+    ],
+    "blocks.10.block.attn.canon._fallback.conv.weight": [
+      384,
+      1,
+      4
+    ],
+    "blocks.10.block.attn.canon.conv.bias": [
+      384
+    ],
+    "blocks.10.block.attn.canon.conv.weight": [
+      384,
+      1,
+      4,
+      4
+    ],
+    "blocks.10.block.attn.module.proj.0.bias": [
+      384
+    ],
+    "blocks.10.block.attn.module.proj.0.weight": [
+      384,
+      384
+    ],
+    "blocks.10.block.attn.module.proj.1._fallback.conv.bias": [
+      384
+    ],
+    "blocks.10.block.attn.module.proj.1._fallback.conv.weight": [
+      384,
+      1,
+      4
+    ],
+    "blocks.10.block.attn.module.proj.1.conv.bias": [
+      384
+    ],
+    "blocks.10.block.attn.module.proj.1.conv.weight": [
+      384,
+      1,
+      4,
+      4
+    ],
+    "blocks.10.block.attn.module.qkv.bias": [
+      1152
+    ],
+    "blocks.10.block.attn.module.qkv.weight": [
+      1152,
+      384
+    ],
+    "blocks.10.block.mlp.canon._fallback.conv.bias": [
+      384
+    ],
+    "blocks.10.block.mlp.canon._fallback.conv.weight": [
+      384,
+      1,
+      4
+    ],
+    "blocks.10.block.mlp.canon.conv.bias": [
+      384
+    ],
+    "blocks.10.block.mlp.canon.conv.weight": [
+      384,
+      1,
+      4,
+      4
+    ],
+    "blocks.10.block.mlp.module.fc1.canon._fallback.conv.bias": [
+      1536
+    ],
+    "blocks.10.block.mlp.module.fc1.canon._fallback.conv.weight": [
+      1536,
+      1,
+      4
+    ],
+    "blocks.10.block.mlp.module.fc1.canon.conv.bias": [
+      1536
+    ],
+    "blocks.10.block.mlp.module.fc1.canon.conv.weight": [
+      1536,
+      1,
+      4,
+      4
+    ],
+    "blocks.10.block.mlp.module.fc1.fc1.bias": [
+      1536
+    ],
+    "blocks.10.block.mlp.module.fc1.fc1.weight": [
+      1536,
+      384
+    ],
+    "blocks.10.block.mlp.module.fc2.bias": [
+      384
+    ],
+    "blocks.10.block.mlp.module.fc2.weight": [
+      384,
+      1536
+    ],
+    "blocks.10.block.norm1.bias": [
+      384
+    ],
+    "blocks.10.block.norm1.weight": [
+      384
+    ],
+    "blocks.10.block.norm2.bias": [
+      384
+    ],
+    "blocks.10.block.norm2.weight": [
+      384
+    ],
+    "blocks.11.block.attn.canon._fallback.conv.bias": [
+      384
+    ],
+    "blocks.11.block.attn.canon._fallback.conv.weight": [
+      384,
+      1,
+      4
+    ],
+    "blocks.11.block.attn.canon.conv.bias": [
+      384
+    ],
+    "blocks.11.block.attn.canon.conv.weight": [
+      384,
+      1,
+      4,
+      4
+    ],
+    "blocks.11.block.attn.module.proj.0.bias": [
+      384
+    ],
+    "blocks.11.block.attn.module.proj.0.weight": [
+      384,
+      384
+    ],
+    "blocks.11.block.attn.module.proj.1._fallback.conv.bias": [
+      384
+    ],
+    "blocks.11.block.attn.module.proj.1._fallback.conv.weight": [
+      384,
+      1,
+      4
+    ],
+    "blocks.11.block.attn.module.proj.1.conv.bias": [
+      384
+    ],
+    "blocks.11.block.attn.module.proj.1.conv.weight": [
+      384,
+      1,
+      4,
+      4
+    ],
+    "blocks.11.block.attn.module.qkv.bias": [
+      1152
+    ],
+    "blocks.11.block.attn.module.qkv.weight": [
+      1152,
+      384
+    ],
+    "blocks.11.block.mlp.canon._fallback.conv.bias": [
+      384
+    ],
+    "blocks.11.block.mlp.canon._fallback.conv.weight": [
+      384,
+      1,
+      4
+    ],
+    "blocks.11.block.mlp.canon.conv.bias": [
+      384
+    ],
+    "blocks.11.block.mlp.canon.conv.weight": [
+      384,
+      1,
+      4,
+      4
+    ],
+    "blocks.11.block.mlp.module.fc1.canon._fallback.conv.bias": [
+      1536
+    ],
+    "blocks.11.block.mlp.module.fc1.canon._fallback.conv.weight": [
+      1536,
+      1,
+      4
+    ],
+    "blocks.11.block.mlp.module.fc1.canon.conv.bias": [
+      1536
+    ],
+    "blocks.11.block.mlp.module.fc1.canon.conv.weight": [
+      1536,
+      1,
+      4,
+      4
+    ],
+    "blocks.11.block.mlp.module.fc1.fc1.bias": [
+      1536
+    ],
+    "blocks.11.block.mlp.module.fc1.fc1.weight": [
+      1536,
+      384
+    ],
+    "blocks.11.block.mlp.module.fc2.bias": [
+      384
+    ],
+    "blocks.11.block.mlp.module.fc2.weight": [
+      384,
+      1536
+    ],
+    "blocks.11.block.norm1.bias": [
+      384
+    ],
+    "blocks.11.block.norm1.weight": [
+      384
+    ],
+    "blocks.11.block.norm2.bias": [
+      384
+    ],
+    "blocks.11.block.norm2.weight": [
+      384
+    ],
+    "blocks.2.block.attn.canon._fallback.conv.bias": [
+      384
+    ],
+    "blocks.2.block.attn.canon._fallback.conv.weight": [
+      384,
+      1,
+      4
+    ],
+    "blocks.2.block.attn.canon.conv.bias": [
+      384
+    ],
+    "blocks.2.block.attn.canon.conv.weight": [
+      384,
+      1,
+      4,
+      4
+    ],
+    "blocks.2.block.attn.module.proj.0.bias": [
+      384
+    ],
+    "blocks.2.block.attn.module.proj.0.weight": [
+      384,
+      384
+    ],
+    "blocks.2.block.attn.module.proj.1._fallback.conv.bias": [
+      384
+    ],
+    "blocks.2.block.attn.module.proj.1._fallback.conv.weight": [
+      384,
+      1,
+      4
+    ],
+    "blocks.2.block.attn.module.proj.1.conv.bias": [
+      384
+    ],
+    "blocks.2.block.attn.module.proj.1.conv.weight": [
+      384,
+      1,
+      4,
+      4
+    ],
+    "blocks.2.block.attn.module.qkv.bias": [
+      1152
+    ],
+    "blocks.2.block.attn.module.qkv.weight": [
+      1152,
+      384
+    ],
+    "blocks.2.block.mlp.canon._fallback.conv.bias": [
+      384
+    ],
+    "blocks.2.block.mlp.canon._fallback.conv.weight": [
+      384,
+      1,
+      4
+    ],
+    "blocks.2.block.mlp.canon.conv.bias": [
+      384
+    ],
+    "blocks.2.block.mlp.canon.conv.weight": [
+      384,
+      1,
+      4,
+      4
+    ],
+    "blocks.2.block.mlp.module.fc1.canon._fallback.conv.bias": [
+      1536
+    ],
+    "blocks.2.block.mlp.module.fc1.canon._fallback.conv.weight": [
+      1536,
+      1,
+      4
+    ],
+    "blocks.2.block.mlp.module.fc1.canon.conv.bias": [
+      1536
+    ],
+    "blocks.2.block.mlp.module.fc1.canon.conv.weight": [
+      1536,
+      1,
+      4,
+      4
+    ],
+    "blocks.2.block.mlp.module.fc1.fc1.bias": [
+      1536
+    ],
+    "blocks.2.block.mlp.module.fc1.fc1.weight": [
+      1536,
+      384
+    ],
+    "blocks.2.block.mlp.module.fc2.bias": [
+      384
+    ],
+    "blocks.2.block.mlp.module.fc2.weight": [
+      384,
+      1536
+    ],
+    "blocks.2.block.norm1.bias": [
+      384
+    ],
+    "blocks.2.block.norm1.weight": [
+      384
+    ],
+    "blocks.2.block.norm2.bias": [
+      384
+    ],
+    "blocks.2.block.norm2.weight": [
+      384
+    ],
+    "blocks.3.block.attn.canon._fallback.conv.bias": [
+      384
+    ],
+    "blocks.3.block.attn.canon._fallback.conv.weight": [
+      384,
+      1,
+      4
+    ],
+    "blocks.3.block.attn.canon.conv.bias": [
+      384
+    ],
+    "blocks.3.block.attn.canon.conv.weight": [
+      384,
+      1,
+      4,
+      4
+    ],
+    "blocks.3.block.attn.module.proj.0.bias": [
+      384
+    ],
+    "blocks.3.block.attn.module.proj.0.weight": [
+      384,
+      384
+    ],
+    "blocks.3.block.attn.module.proj.1._fallback.conv.bias": [
+      384
+    ],
+    "blocks.3.block.attn.module.proj.1._fallback.conv.weight": [
+      384,
+      1,
+      4
+    ],
+    "blocks.3.block.attn.module.proj.1.conv.bias": [
+      384
+    ],
+    "blocks.3.block.attn.module.proj.1.conv.weight": [
+      384,
+      1,
+      4,
+      4
+    ],
+    "blocks.3.block.attn.module.qkv.bias": [
+      1152
+    ],
+    "blocks.3.block.attn.module.qkv.weight": [
+      1152,
+      384
+    ],
+    "blocks.3.block.mlp.canon._fallback.conv.bias": [
+      384
+    ],
+    "blocks.3.block.mlp.canon._fallback.conv.weight": [
+      384,
+      1,
+      4
+    ],
+    "blocks.3.block.mlp.canon.conv.bias": [
+      384
+    ],
+    "blocks.3.block.mlp.canon.conv.weight": [
+      384,
+      1,
+      4,
+      4
+    ],
+    "blocks.3.block.mlp.module.fc1.canon._fallback.conv.bias": [
+      1536
+    ],
+    "blocks.3.block.mlp.module.fc1.canon._fallback.conv.weight": [
+      1536,
+      1,
+      4
+    ],
+    "blocks.3.block.mlp.module.fc1.canon.conv.bias": [
+      1536
+    ],
+    "blocks.3.block.mlp.module.fc1.canon.conv.weight": [
+      1536,
+      1,
+      4,
+      4
+    ],
+    "blocks.3.block.mlp.module.fc1.fc1.bias": [
+      1536
+    ],
+    "blocks.3.block.mlp.module.fc1.fc1.weight": [
+      1536,
+      384
+    ],
+    "blocks.3.block.mlp.module.fc2.bias": [
+      384
+    ],
+    "blocks.3.block.mlp.module.fc2.weight": [
+      384,
+      1536
+    ],
+    "blocks.3.block.norm1.bias": [
+      384
+    ],
+    "blocks.3.block.norm1.weight": [
+      384
+    ],
+    "blocks.3.block.norm2.bias": [
+      384
+    ],
+    "blocks.3.block.norm2.weight": [
+      384
+    ],
+    "blocks.4.block.attn.canon._fallback.conv.bias": [
+      384
+    ],
+    "blocks.4.block.attn.canon._fallback.conv.weight": [
+      384,
+      1,
+      4
+    ],
+    "blocks.4.block.attn.canon.conv.bias": [
+      384
+    ],
+    "blocks.4.block.attn.canon.conv.weight": [
+      384,
+      1,
+      4,
+      4
+    ],
+    "blocks.4.block.attn.module.proj.0.bias": [
+      384
+    ],
+    "blocks.4.block.attn.module.proj.0.weight": [
+      384,
+      384
+    ],
+    "blocks.4.block.attn.module.proj.1._fallback.conv.bias": [
+      384
+    ],
+    "blocks.4.block.attn.module.proj.1._fallback.conv.weight": [
+      384,
+      1,
+      4
+    ],
+    "blocks.4.block.attn.module.proj.1.conv.bias": [
+      384
+    ],
+    "blocks.4.block.attn.module.proj.1.conv.weight": [
+      384,
+      1,
+      4,
+      4
+    ],
+    "blocks.4.block.attn.module.qkv.bias": [
+      1152
+    ],
+    "blocks.4.block.attn.module.qkv.weight": [
+      1152,
+      384
+    ],
+    "blocks.4.block.mlp.canon._fallback.conv.bias": [
+      384
+    ],
+    "blocks.4.block.mlp.canon._fallback.conv.weight": [
+      384,
+      1,
+      4
+    ],
+    "blocks.4.block.mlp.canon.conv.bias": [
+      384
+    ],
+    "blocks.4.block.mlp.canon.conv.weight": [
+      384,
+      1,
+      4,
+      4
+    ],
+    "blocks.4.block.mlp.module.fc1.canon._fallback.conv.bias": [
+      1536
+    ],
+    "blocks.4.block.mlp.module.fc1.canon._fallback.conv.weight": [
+      1536,
+      1,
+      4
+    ],
+    "blocks.4.block.mlp.module.fc1.canon.conv.bias": [
+      1536
+    ],
+    "blocks.4.block.mlp.module.fc1.canon.conv.weight": [
+      1536,
+      1,
+      4,
+      4
+    ],
+    "blocks.4.block.mlp.module.fc1.fc1.bias": [
+      1536
+    ],
+    "blocks.4.block.mlp.module.fc1.fc1.weight": [
+      1536,
+      384
+    ],
+    "blocks.4.block.mlp.module.fc2.bias": [
+      384
+    ],
+    "blocks.4.block.mlp.module.fc2.weight": [
+      384,
+      1536
+    ],
+    "blocks.4.block.norm1.bias": [
+      384
+    ],
+    "blocks.4.block.norm1.weight": [
+      384
+    ],
+    "blocks.4.block.norm2.bias": [
+      384
+    ],
+    "blocks.4.block.norm2.weight": [
+      384
+    ],
+    "blocks.5.block.attn.canon._fallback.conv.bias": [
+      384
+    ],
+    "blocks.5.block.attn.canon._fallback.conv.weight": [
+      384,
+      1,
+      4
+    ],
+    "blocks.5.block.attn.canon.conv.bias": [
+      384
+    ],
+    "blocks.5.block.attn.canon.conv.weight": [
+      384,
+      1,
+      4,
+      4
+    ],
+    "blocks.5.block.attn.module.proj.0.bias": [
+      384
+    ],
+    "blocks.5.block.attn.module.proj.0.weight": [
+      384,
+      384
+    ],
+    "blocks.5.block.attn.module.proj.1._fallback.conv.bias": [
+      384
+    ],
+    "blocks.5.block.attn.module.proj.1._fallback.conv.weight": [
+      384,
+      1,
+      4
+    ],
+    "blocks.5.block.attn.module.proj.1.conv.bias": [
+      384
+    ],
+    "blocks.5.block.attn.module.proj.1.conv.weight": [
+      384,
+      1,
+      4,
+      4
+    ],
+    "blocks.5.block.attn.module.qkv.bias": [
+      1152
+    ],
+    "blocks.5.block.attn.module.qkv.weight": [
+      1152,
+      384
+    ],
+    "blocks.5.block.mlp.canon._fallback.conv.bias": [
+      384
+    ],
+    "blocks.5.block.mlp.canon._fallback.conv.weight": [
+      384,
+      1,
+      4
+    ],
+    "blocks.5.block.mlp.canon.conv.bias": [
+      384
+    ],
+    "blocks.5.block.mlp.canon.conv.weight": [
+      384,
+      1,
+      4,
+      4
+    ],
+    "blocks.5.block.mlp.module.fc1.canon._fallback.conv.bias": [
+      1536
+    ],
+    "blocks.5.block.mlp.module.fc1.canon._fallback.conv.weight": [
+      1536,
+      1,
+      4
+    ],
+    "blocks.5.block.mlp.module.fc1.canon.conv.bias": [
+      1536
+    ],
+    "blocks.5.block.mlp.module.fc1.canon.conv.weight": [
+      1536,
+      1,
+      4,
+      4
+    ],
+    "blocks.5.block.mlp.module.fc1.fc1.bias": [
+      1536
+    ],
+    "blocks.5.block.mlp.module.fc1.fc1.weight": [
+      1536,
+      384
+    ],
+    "blocks.5.block.mlp.module.fc2.bias": [
+      384
+    ],
+    "blocks.5.block.mlp.module.fc2.weight": [
+      384,
+      1536
+    ],
+    "blocks.5.block.norm1.bias": [
+      384
+    ],
+    "blocks.5.block.norm1.weight": [
+      384
+    ],
+    "blocks.5.block.norm2.bias": [
+      384
+    ],
+    "blocks.5.block.norm2.weight": [
+      384
+    ],
+    "blocks.6.block.attn.canon._fallback.conv.bias": [
+      384
+    ],
+    "blocks.6.block.attn.canon._fallback.conv.weight": [
+      384,
+      1,
+      4
+    ],
+    "blocks.6.block.attn.canon.conv.bias": [
+      384
+    ],
+    "blocks.6.block.attn.canon.conv.weight": [
+      384,
+      1,
+      4,
+      4
+    ],
+    "blocks.6.block.attn.module.proj.0.bias": [
+      384
+    ],
+    "blocks.6.block.attn.module.proj.0.weight": [
+      384,
+      384
+    ],
+    "blocks.6.block.attn.module.proj.1._fallback.conv.bias": [
+      384
+    ],
+    "blocks.6.block.attn.module.proj.1._fallback.conv.weight": [
+      384,
+      1,
+      4
+    ],
+    "blocks.6.block.attn.module.proj.1.conv.bias": [
+      384
+    ],
+    "blocks.6.block.attn.module.proj.1.conv.weight": [
+      384,
+      1,
+      4,
+      4
+    ],
+    "blocks.6.block.attn.module.qkv.bias": [
+      1152
+    ],
+    "blocks.6.block.attn.module.qkv.weight": [
+      1152,
+      384
+    ],
+    "blocks.6.block.mlp.canon._fallback.conv.bias": [
+      384
+    ],
+    "blocks.6.block.mlp.canon._fallback.conv.weight": [
+      384,
+      1,
+      4
+    ],
+    "blocks.6.block.mlp.canon.conv.bias": [
+      384
+    ],
+    "blocks.6.block.mlp.canon.conv.weight": [
+      384,
+      1,
+      4,
+      4
+    ],
+    "blocks.6.block.mlp.module.fc1.canon._fallback.conv.bias": [
+      1536
+    ],
+    "blocks.6.block.mlp.module.fc1.canon._fallback.conv.weight": [
+      1536,
+      1,
+      4
+    ],
+    "blocks.6.block.mlp.module.fc1.canon.conv.bias": [
+      1536
+    ],
+    "blocks.6.block.mlp.module.fc1.canon.conv.weight": [
+      1536,
+      1,
+      4,
+      4
+    ],
+    "blocks.6.block.mlp.module.fc1.fc1.bias": [
+      1536
+    ],
+    "blocks.6.block.mlp.module.fc1.fc1.weight": [
+      1536,
+      384
+    ],
+    "blocks.6.block.mlp.module.fc2.bias": [
+      384
+    ],
+    "blocks.6.block.mlp.module.fc2.weight": [
+      384,
+      1536
+    ],
+    "blocks.6.block.norm1.bias": [
+      384
+    ],
+    "blocks.6.block.norm1.weight": [
+      384
+    ],
+    "blocks.6.block.norm2.bias": [
+      384
+    ],
+    "blocks.6.block.norm2.weight": [
+      384
+    ],
+    "blocks.7.block.attn.canon._fallback.conv.bias": [
+      384
+    ],
+    "blocks.7.block.attn.canon._fallback.conv.weight": [
+      384,
+      1,
+      4
+    ],
+    "blocks.7.block.attn.canon.conv.bias": [
+      384
+    ],
+    "blocks.7.block.attn.canon.conv.weight": [
+      384,
+      1,
+      4,
+      4
+    ],
+    "blocks.7.block.attn.module.proj.0.bias": [
+      384
+    ],
+    "blocks.7.block.attn.module.proj.0.weight": [
+      384,
+      384
+    ],
+    "blocks.7.block.attn.module.proj.1._fallback.conv.bias": [
+      384
+    ],
+    "blocks.7.block.attn.module.proj.1._fallback.conv.weight": [
+      384,
+      1,
+      4
+    ],
+    "blocks.7.block.attn.module.proj.1.conv.bias": [
+      384
+    ],
+    "blocks.7.block.attn.module.proj.1.conv.weight": [
+      384,
+      1,
+      4,
+      4
+    ],
+    "blocks.7.block.attn.module.qkv.bias": [
+      1152
+    ],
+    "blocks.7.block.attn.module.qkv.weight": [
+      1152,
+      384
+    ],
+    "blocks.7.block.mlp.canon._fallback.conv.bias": [
+      384
+    ],
+    "blocks.7.block.mlp.canon._fallback.conv.weight": [
+      384,
+      1,
+      4
+    ],
+    "blocks.7.block.mlp.canon.conv.bias": [
+      384
+    ],
+    "blocks.7.block.mlp.canon.conv.weight": [
+      384,
+      1,
+      4,
+      4
+    ],
+    "blocks.7.block.mlp.module.fc1.canon._fallback.conv.bias": [
+      1536
+    ],
+    "blocks.7.block.mlp.module.fc1.canon._fallback.conv.weight": [
+      1536,
+      1,
+      4
+    ],
+    "blocks.7.block.mlp.module.fc1.canon.conv.bias": [
+      1536
+    ],
+    "blocks.7.block.mlp.module.fc1.canon.conv.weight": [
+      1536,
+      1,
+      4,
+      4
+    ],
+    "blocks.7.block.mlp.module.fc1.fc1.bias": [
+      1536
+    ],
+    "blocks.7.block.mlp.module.fc1.fc1.weight": [
+      1536,
+      384
+    ],
+    "blocks.7.block.mlp.module.fc2.bias": [
+      384
+    ],
+    "blocks.7.block.mlp.module.fc2.weight": [
+      384,
+      1536
+    ],
+    "blocks.7.block.norm1.bias": [
+      384
+    ],
+    "blocks.7.block.norm1.weight": [
+      384
+    ],
+    "blocks.7.block.norm2.bias": [
+      384
+    ],
+    "blocks.7.block.norm2.weight": [
+      384
+    ],
+    "blocks.8.block.attn.canon._fallback.conv.bias": [
+      384
+    ],
+    "blocks.8.block.attn.canon._fallback.conv.weight": [
+      384,
+      1,
+      4
+    ],
+    "blocks.8.block.attn.canon.conv.bias": [
+      384
+    ],
+    "blocks.8.block.attn.canon.conv.weight": [
+      384,
+      1,
+      4,
+      4
+    ],
+    "blocks.8.block.attn.module.proj.0.bias": [
+      384
+    ],
+    "blocks.8.block.attn.module.proj.0.weight": [
+      384,
+      384
+    ],
+    "blocks.8.block.attn.module.proj.1._fallback.conv.bias": [
+      384
+    ],
+    "blocks.8.block.attn.module.proj.1._fallback.conv.weight": [
+      384,
+      1,
+      4
+    ],
+    "blocks.8.block.attn.module.proj.1.conv.bias": [
+      384
+    ],
+    "blocks.8.block.attn.module.proj.1.conv.weight": [
+      384,
+      1,
+      4,
+      4
+    ],
+    "blocks.8.block.attn.module.qkv.bias": [
+      1152
+    ],
+    "blocks.8.block.attn.module.qkv.weight": [
+      1152,
+      384
+    ],
+    "blocks.8.block.mlp.canon._fallback.conv.bias": [
+      384
+    ],
+    "blocks.8.block.mlp.canon._fallback.conv.weight": [
+      384,
+      1,
+      4
+    ],
+    "blocks.8.block.mlp.canon.conv.bias": [
+      384
+    ],
+    "blocks.8.block.mlp.canon.conv.weight": [
+      384,
+      1,
+      4,
+      4
+    ],
+    "blocks.8.block.mlp.module.fc1.canon._fallback.conv.bias": [
+      1536
+    ],
+    "blocks.8.block.mlp.module.fc1.canon._fallback.conv.weight": [
+      1536,
+      1,
+      4
+    ],
+    "blocks.8.block.mlp.module.fc1.canon.conv.bias": [
+      1536
+    ],
+    "blocks.8.block.mlp.module.fc1.canon.conv.weight": [
+      1536,
+      1,
+      4,
+      4
+    ],
+    "blocks.8.block.mlp.module.fc1.fc1.bias": [
+      1536
+    ],
+    "blocks.8.block.mlp.module.fc1.fc1.weight": [
+      1536,
+      384
+    ],
+    "blocks.8.block.mlp.module.fc2.bias": [
+      384
+    ],
+    "blocks.8.block.mlp.module.fc2.weight": [
+      384,
+      1536
+    ],
+    "blocks.8.block.norm1.bias": [
+      384
+    ],
+    "blocks.8.block.norm1.weight": [
+      384
+    ],
+    "blocks.8.block.norm2.bias": [
+      384
+    ],
+    "blocks.8.block.norm2.weight": [
+      384
+    ],
+    "blocks.9.block.attn.canon._fallback.conv.bias": [
+      384
+    ],
+    "blocks.9.block.attn.canon._fallback.conv.weight": [
+      384,
+      1,
+      4
+    ],
+    "blocks.9.block.attn.canon.conv.bias": [
+      384
+    ],
+    "blocks.9.block.attn.canon.conv.weight": [
+      384,
+      1,
+      4,
+      4
+    ],
+    "blocks.9.block.attn.module.proj.0.bias": [
+      384
+    ],
+    "blocks.9.block.attn.module.proj.0.weight": [
+      384,
+      384
+    ],
+    "blocks.9.block.attn.module.proj.1._fallback.conv.bias": [
+      384
+    ],
+    "blocks.9.block.attn.module.proj.1._fallback.conv.weight": [
+      384,
+      1,
+      4
+    ],
+    "blocks.9.block.attn.module.proj.1.conv.bias": [
+      384
+    ],
+    "blocks.9.block.attn.module.proj.1.conv.weight": [
+      384,
+      1,
+      4,
+      4
+    ],
+    "blocks.9.block.attn.module.qkv.bias": [
+      1152
+    ],
+    "blocks.9.block.attn.module.qkv.weight": [
+      1152,
+      384
+    ],
+    "blocks.9.block.mlp.canon._fallback.conv.bias": [
+      384
+    ],
+    "blocks.9.block.mlp.canon._fallback.conv.weight": [
+      384,
+      1,
+      4
+    ],
+    "blocks.9.block.mlp.canon.conv.bias": [
+      384
+    ],
+    "blocks.9.block.mlp.canon.conv.weight": [
+      384,
+      1,
+      4,
+      4
+    ],
+    "blocks.9.block.mlp.module.fc1.canon._fallback.conv.bias": [
+      1536
+    ],
+    "blocks.9.block.mlp.module.fc1.canon._fallback.conv.weight": [
+      1536,
+      1,
+      4
+    ],
+    "blocks.9.block.mlp.module.fc1.canon.conv.bias": [
+      1536
+    ],
+    "blocks.9.block.mlp.module.fc1.canon.conv.weight": [
+      1536,
+      1,
+      4,
+      4
+    ],
+    "blocks.9.block.mlp.module.fc1.fc1.bias": [
+      1536
+    ],
+    "blocks.9.block.mlp.module.fc1.fc1.weight": [
+      1536,
+      384
+    ],
+    "blocks.9.block.mlp.module.fc2.bias": [
+      384
+    ],
+    "blocks.9.block.mlp.module.fc2.weight": [
+      384,
+      1536
+    ],
+    "blocks.9.block.norm1.bias": [
+      384
+    ],
+    "blocks.9.block.norm1.weight": [
+      384
+    ],
+    "blocks.9.block.norm2.bias": [
+      384
+    ],
+    "blocks.9.block.norm2.weight": [
+      384
+    ],
+    "cls_token": [
+      1,
+      1,
+      384
+    ],
+    "fc_norm.bias": [
+      384
+    ],
+    "fc_norm.weight": [
+      384
+    ],
+    "patch_embed.proj.bias": [
+      384
+    ],
+    "patch_embed.proj.weight": [
+      384,
+      1,
+      16,
+      16
+    ],
+    "pos_embed": [
+      1,
+      97,
+      384
+    ]
+  },
+  "student_state_dict_size": 342,
+  "total_parameters": 22337408
+}

modeling_hear_canon.py

@@ -0,0 +1,707 @@
+"""PyTorch model definition for the distilled HeAR ViT-S + Canon model."""
+
+from __future__ import annotations
+
+import math
+from typing import Optional, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import PreTrainedModel
+from transformers.modeling_outputs import BaseModelOutputWithPooling
+
+try:
+  from .configuration_hear_canon import HearCanonViTConfig
+except ImportError:  # pragma: no cover
+  from configuration_hear_canon import HearCanonViTConfig
+
+
+_MEL_WEIGHT_CACHE: dict[tuple, torch.Tensor] = {}
+_WINDOW_CACHE: dict[tuple, torch.Tensor] = {}
+
+
+def _enclosing_power_of_two(value: int) -> int:
+  return int(2 ** math.ceil(math.log2(value))) if value > 0 else 1
+
+
+def _compute_stft(
+  signals: torch.Tensor,
+  frame_length: int,
+  frame_step: int,
+  fft_length: int | None = None,
+  window_fn=torch.hann_window,
+  pad_end: bool = True,
+) -> torch.Tensor:
+  if signals.ndim < 1:
+    raise ValueError(f"Input signals must have rank at least 1, got rank {signals.ndim}.")
+  if fft_length is None:
+    fft_length = _enclosing_power_of_two(frame_length)
+
+  if pad_end:
+    n_frames = math.ceil(signals.shape[-1] / frame_step) if signals.shape[-1] > 0 else 0
+    padded_length = max(0, (n_frames - 1) * frame_step + frame_length) if n_frames > 0 else frame_length
+    padding_needed = max(0, padded_length - signals.shape[-1])
+    if padding_needed > 0:
+      signals = F.pad(signals, (0, padding_needed))
+
+  framed_signals = signals.unfold(-1, frame_length, frame_step)
+  if framed_signals.shape[-2] == 0:
+    return torch.empty(
+      *signals.shape[:-1],
+      0,
+      fft_length // 2 + 1,
+      dtype=torch.complex64,
+      device=signals.device,
+    )
+
+  if window_fn is not None:
+    if window_fn is torch.hann_window:
+      key = (str(framed_signals.device), framed_signals.dtype, int(frame_length))
+      window = _WINDOW_CACHE.get(key)
+      if window is None:
+        window = torch.hann_window(frame_length, device=framed_signals.device, dtype=framed_signals.dtype)
+        _WINDOW_CACHE[key] = window
+    else:
+      window = window_fn(frame_length).to(framed_signals.device).to(framed_signals.dtype)
+    framed_signals = framed_signals * window
+
+  return torch.fft.rfft(framed_signals, n=fft_length, dim=-1)
+
+
+def _ema(
+  inputs: torch.Tensor,
+  num_channels: int,
+  smooth_coef: float,
+  initial_state: torch.Tensor | None = None,
+) -> torch.Tensor:
+  batch_size, timesteps, _ = inputs.shape
+
+  if initial_state is None:
+    ema_state = torch.zeros((batch_size, num_channels), dtype=inputs.dtype, device=inputs.device)
+  else:
+    ema_state = initial_state
+
+  gain_in = float(smooth_coef)
+  gain_rec = float(1.0 - smooth_coef)
+  output_sequence: list[torch.Tensor] = []
+
+  start = 1 if initial_state is not None else 0
+  if start:
+    output_sequence.append(ema_state)
+  for t in range(start, timesteps):
+    current_input = inputs[:, t, :]
+    ema_state = (current_input * gain_in) + (ema_state * gain_rec)
+    output_sequence.append(ema_state)
+
+  return torch.stack(output_sequence, dim=1)
+
+
+def _pcen_function(
+  inputs: torch.Tensor,
+  num_channels: int = 128,
+  alpha: float = 0.8,
+  smooth_coef: float = 0.04,
+  delta: float = 2.0,
+  root: float = 2.0,
+  floor: float = 1e-8,
+) -> torch.Tensor:
+  alpha_param = torch.full((num_channels,), float(alpha), device=inputs.device, dtype=inputs.dtype)
+  delta_param = torch.full((num_channels,), float(delta), device=inputs.device, dtype=inputs.dtype)
+  root_param = torch.full((num_channels,), float(root), device=inputs.device, dtype=inputs.dtype)
+
+  alpha_param = torch.minimum(alpha_param, torch.ones_like(alpha_param))
+  root_param = torch.maximum(root_param, torch.ones_like(root_param))
+  ema_smoother = _ema(
+    inputs,
+    num_channels=num_channels,
+    smooth_coef=smooth_coef,
+    initial_state=inputs[:, 0] if inputs.ndim > 1 else None,
+  )
+
+  one_over_root = 1.0 / root_param
+  output = (inputs / (floor + ema_smoother) ** alpha_param + delta_param) ** one_over_root - delta_param**one_over_root
+  return output
+
+
+def _hertz_to_mel(frequencies_hertz: torch.Tensor) -> torch.Tensor:
+  return 2595.0 * torch.log10(1.0 + frequencies_hertz / 700.0)
+
+
+def _linear_to_mel_weight_matrix(
+  device: torch.device,
+  num_mel_bins: int = 128,
+  num_spectrogram_bins: int = 201,
+  sample_rate: float = 16000,
+  lower_edge_hertz: float = 0.0,
+  upper_edge_hertz: float = 8000.0,
+  dtype: torch.dtype = torch.float32,
+) -> torch.Tensor:
+  cache_key = (
+    str(device),
+    dtype,
+    int(num_mel_bins),
+    int(num_spectrogram_bins),
+    float(sample_rate),
+    float(lower_edge_hertz),
+    float(upper_edge_hertz),
+  )
+  cached = _MEL_WEIGHT_CACHE.get(cache_key)
+  if cached is not None:
+    return cached
+
+  sample_rate_tensor = torch.tensor(sample_rate, dtype=dtype)
+  lower_edge_hertz_tensor = torch.tensor(lower_edge_hertz, dtype=dtype, device=device)
+  upper_edge_hertz_tensor = torch.tensor(upper_edge_hertz, dtype=dtype, device=device)
+  zero = torch.tensor(0.0, dtype=dtype, device=device)
+
+  bands_to_zero = 1
+  nyquist_hertz = sample_rate_tensor / 2.0
+  linear_frequencies = torch.linspace(zero, nyquist_hertz, num_spectrogram_bins, dtype=dtype, device=device)[bands_to_zero:]
+  spectrogram_bins_mel = _hertz_to_mel(linear_frequencies).unsqueeze(1)
+
+  band_edges_mel = torch.linspace(
+    _hertz_to_mel(lower_edge_hertz_tensor),
+    _hertz_to_mel(upper_edge_hertz_tensor),
+    num_mel_bins + 2,
+    dtype=dtype,
+    device=device,
+  )
+  band_edges_mel = band_edges_mel.unfold(0, 3, 1)
+
+  lower_edge_mel = band_edges_mel[:, 0].unsqueeze(0)
+  center_mel = band_edges_mel[:, 1].unsqueeze(0)
+  upper_edge_mel = band_edges_mel[:, 2].unsqueeze(0)
+
+  lower_slopes = (spectrogram_bins_mel - lower_edge_mel) / (center_mel - lower_edge_mel)
+  upper_slopes = (upper_edge_mel - spectrogram_bins_mel) / (upper_edge_mel - center_mel)
+
+  mel_weights_matrix = torch.maximum(zero, torch.minimum(lower_slopes, upper_slopes))
+  mel_weights_matrix = F.pad(mel_weights_matrix, (0, 0, bands_to_zero, 0), mode="constant", value=0.0)
+
+  _MEL_WEIGHT_CACHE[cache_key] = mel_weights_matrix
+  return mel_weights_matrix
+
+
+def _mel_pcen(x: torch.Tensor) -> torch.Tensor:
+  x = x.float()
+  x = x - torch.min(x)
+  x = x / (torch.max(x) + 1e-8)
+  x = (x * 2) - 1
+
+  frame_length = 16 * 25
+  frame_step = 160
+
+  stft = _compute_stft(
+    x,
+    frame_length=frame_length,
+    fft_length=frame_length,
+    frame_step=frame_step,
+    window_fn=torch.hann_window,
+    pad_end=True,
+  )
+  spectrograms = torch.square(torch.abs(stft))
+
+  mel_transform = _linear_to_mel_weight_matrix(x.device)
+  mel_spectrograms = torch.matmul(spectrograms, mel_transform)
+  return _pcen_function(mel_spectrograms)
+
+
+def _torch_resize_bilinear_tf_compat(images: torch.Tensor, size: Tuple[int, int]) -> torch.Tensor:
+  original_dims = images.dim()
+  new_height, new_width = size
+  if original_dims not in (3, 4):
+    raise ValueError("Input tensor must be 3D [C, H, W] or 4D [B, C, H, W].")
+
+  images = images.to(torch.float32)
+  was_3d = original_dims == 3
+  if was_3d:
+    images = images.unsqueeze(0)
+
+  resized = F.interpolate(
+    images,
+    size=(new_height, new_width),
+    mode="bilinear",
+    align_corners=False,
+    antialias=False,
+  )
+
+  if was_3d:
+    resized = resized.squeeze(0)
+  return resized
+
+
+def preprocess_audio(
+  audio: torch.Tensor,
+  *,
+  clip_samples: int = 32000,
+  target_size: Tuple[int, int] = (192, 128),
+) -> torch.Tensor:
+  """Convert raw 16 kHz waveforms to model-ready mel/PCEN spectrograms."""
+  if audio.ndim == 1:
+    audio = audio.unsqueeze(0)
+  if audio.ndim != 2:
+    raise ValueError(f"Input audio must have rank 2 [B, samples], got rank {audio.ndim}.")
+
+  if audio.shape[1] < clip_samples:
+    n = clip_samples - audio.shape[1]
+    audio = F.pad(audio, pad=(0, n), mode="constant", value=0)
+  elif audio.shape[1] > clip_samples:
+    audio = audio[:, :clip_samples]
+
+  spectrogram = _mel_pcen(audio)
+  spectrogram = torch.unsqueeze(spectrogram, dim=1)
+  return _torch_resize_bilinear_tf_compat(spectrogram, size=target_size)
+
+
+def _disable_positional_embeddings(model: nn.Module) -> None:
+  if not hasattr(model, "pos_embed"):
+    return
+  pos = getattr(model, "pos_embed")
+  if pos is None:
+    return
+  if isinstance(pos, nn.Parameter):
+    with torch.no_grad():
+      new_pos = torch.zeros_like(pos)
+    model.pos_embed = nn.Parameter(new_pos, requires_grad=False)
+    return
+  if torch.is_tensor(pos):
+    with torch.no_grad():
+      pos.zero_()
+
+
+class CanonLayer(nn.Module):
+  def __init__(self, dim: int, kernel_size: int = 4, causal: bool = False) -> None:
+    super().__init__()
+    self.kernel_size = int(kernel_size)
+    self.causal = bool(causal)
+    self.conv = nn.Conv1d(dim, dim, kernel_size=self.kernel_size, groups=dim, bias=True)
+
+  def forward(self, x: torch.Tensor) -> torch.Tensor:
+    y = x.transpose(1, 2)
+    if self.causal:
+      pad_left = self.kernel_size - 1
+      pad_right = 0
+    else:
+      pad_left = (self.kernel_size - 1) // 2
+      pad_right = self.kernel_size // 2
+    y = F.pad(y, (pad_left, pad_right))
+    y = self.conv(y)
+    y = y.transpose(1, 2)
+    return x + y
+
+
+class Canon2DLayer(nn.Module):
+  def __init__(self, dim: int, kernel_h: int, kernel_w: int, causal_time: bool = False) -> None:
+    super().__init__()
+    self.kernel_h = int(kernel_h)
+    self.kernel_w = int(kernel_w)
+    self.causal_time = bool(causal_time)
+    self.conv = nn.Conv2d(dim, dim, kernel_size=(self.kernel_h, self.kernel_w), groups=dim, bias=True)
+    self.grid_size: Optional[Tuple[int, int]] = None
+    self.expect_cls: Optional[bool] = None
+    self._warned = False
+    self._fallback = CanonLayer(dim, kernel_size=self.kernel_h, causal=self.causal_time)
+
+  def _warn_once(self, msg: str) -> None:
+    if not self._warned:
+      print(msg, flush=True)
+      self._warned = True
+
+  def forward(self, x: torch.Tensor) -> torch.Tensor:
+    if x.ndim != 3:
+      raise RuntimeError("Canon2D expects input of shape [B, N, C].")
+
+    if self.grid_size is None:
+      self._warn_once("Warning: Canon2D missing grid_size; falling back to 1D Canon.")
+      return self._fallback(x)
+
+    h, w = self.grid_size
+    if not (isinstance(h, int) and isinstance(w, int)):
+      self._warn_once("Warning: Canon2D grid_size is invalid; falling back to 1D Canon.")
+      return self._fallback(x)
+
+    b, n, c = x.shape
+    expected = int(h) * int(w)
+    if self.expect_cls is True and n != expected + 1:
+      self._warn_once(
+        f"Warning: Canon2D expected CLS token with N=1+H*W ({expected + 1}) but got N={n}; falling back to 1D Canon."
+      )
+      return self._fallback(x)
+
+    has_cls = False
+    if n == expected + 1:
+      has_cls = True
+      cls = x[:, :1, :]
+      patches = x[:, 1:, :]
+    elif n == expected:
+      cls = None
+      patches = x
+    else:
+      self._warn_once(
+        f"Warning: Canon2D token count mismatch (N={n}, H*W={expected}); falling back to 1D Canon."
+      )
+      return self._fallback(x)
+
+    patches = patches.transpose(1, 2).contiguous().view(b, c, int(h), int(w))
+    if self.causal_time:
+      pad_h_top = self.kernel_h - 1
+      pad_h_bottom = 0
+    else:
+      pad_h_top = (self.kernel_h - 1) // 2
+      pad_h_bottom = self.kernel_h // 2
+    pad_w_left = (self.kernel_w - 1) // 2
+    pad_w_right = self.kernel_w // 2
+    patches = F.pad(patches, (pad_w_left, pad_w_right, pad_h_top, pad_h_bottom))
+    y = self.conv(patches)
+    y = y.view(b, c, int(h) * int(w)).transpose(1, 2)
+    if has_cls:
+      zero_cls = torch.zeros_like(cls)
+      y = torch.cat([zero_cls, y], dim=1)
+    return x + y
+
+
+class CanonInputWrapper(nn.Module):
+  def __init__(self, module: nn.Module, canon: nn.Module) -> None:
+    super().__init__()
+    self.module = module
+    self.canon = canon
+
+  def forward(self, x: torch.Tensor, *args, **kwargs) -> torch.Tensor:
+    return self.module(self.canon(x), *args, **kwargs)
+
+
+class CanonQKVWrapper(nn.Module):
+  def __init__(self, qkv: nn.Module, canon: nn.Module) -> None:
+    super().__init__()
+    self.qkv = qkv
+    self.canon = canon
+
+  def forward(self, x: torch.Tensor, *args, **kwargs) -> torch.Tensor:
+    y = self.qkv(x, *args, **kwargs)
+    if y.ndim != 3:
+      raise RuntimeError("Canon-B expects QKV output of shape [B, N, 3*D].")
+    return self.canon(y)
+
+
+class CanonFC1Wrapper(nn.Module):
+  def __init__(self, fc1: nn.Module, canon: nn.Module) -> None:
+    super().__init__()
+    self.fc1 = fc1
+    self.canon = canon
+
+  def forward(self, x: torch.Tensor, *args, **kwargs) -> torch.Tensor:
+    y = self.fc1(x, *args, **kwargs)
+    if y.ndim != 3:
+      raise RuntimeError("Canon-D expects MLP FC1 output of shape [B, N, M].")
+    return self.canon(y)
+
+
+class CanonBlockWrapper(nn.Module):
+  def __init__(
+    self,
+    block: nn.Module,
+    dim: int,
+    *,
+    kernel_size: int = 4,
+    canon_a: bool = False,
+    canon_b: bool = False,
+    canon_b_qkv: bool = False,
+    canon_c: bool = False,
+    canon_d: bool = False,
+    causal: bool = False,
+    use_2d: bool = False,
+    grid_size: Optional[Tuple[int, int]] = None,
+    expect_cls: Optional[bool] = None,
+  ) -> None:
+    super().__init__()
+    self.block = block
+    self.use_2d = bool(use_2d)
+    self.grid_size = tuple(grid_size) if grid_size is not None else None
+    self.expect_cls = expect_cls
+    self._insert_canon(
+      dim=dim,
+      kernel_size=kernel_size,
+      canon_a=canon_a,
+      canon_b=canon_b,
+      canon_b_qkv=canon_b_qkv,
+      canon_c=canon_c,
+      canon_d=canon_d,
+      causal=causal,
+    )
+
+  def _make_canon(self, dim: int, kernel_size: int, causal: bool) -> nn.Module:
+    if self.use_2d:
+      canon = Canon2DLayer(int(dim), int(kernel_size), int(kernel_size), causal_time=causal)
+      canon.grid_size = self.grid_size
+      canon.expect_cls = self.expect_cls
+      return canon
+    return CanonLayer(int(dim), kernel_size=kernel_size, causal=causal)
+
+  def _insert_canon(
+    self,
+    *,
+    dim: int,
+    kernel_size: int,
+    canon_a: bool,
+    canon_b: bool,
+    canon_b_qkv: bool,
+    canon_c: bool,
+    canon_d: bool,
+    causal: bool,
+  ) -> None:
+    block = self.block
+
+    if canon_b:
+      if not hasattr(block, "attn"):
+        raise RuntimeError("Canon-B requested but block has no `.attn`.")
+      attn = block.attn
+      if canon_b_qkv:
+        if not hasattr(attn, "qkv"):
+          raise RuntimeError("Canon-B(QKV) requested but attention has no `.qkv`.")
+        qkv = attn.qkv
+        qkv_dim = getattr(qkv, "out_features", None)
+        if qkv_dim is None:
+          raise RuntimeError("Canon-B(QKV) requested but could not read qkv out_features.")
+        attn.qkv = CanonQKVWrapper(qkv, self._make_canon(int(qkv_dim), kernel_size, causal))
+      else:
+        if not hasattr(attn, "proj"):
+          raise RuntimeError("Canon-B requested but attention has no `.proj`.")
+        attn.proj = nn.Sequential(attn.proj, self._make_canon(int(dim), kernel_size, causal))
+
+    if canon_a:
+      if not hasattr(block, "attn"):
+        raise RuntimeError("Canon-A requested but block has no `.attn`.")
+      block.attn = CanonInputWrapper(block.attn, self._make_canon(int(dim), kernel_size, causal))
+
+    if canon_d:
+      if not hasattr(block, "mlp"):
+        raise RuntimeError("Canon-D requested but block has no `.mlp`.")
+      mlp = block.mlp
+      if not hasattr(mlp, "fc1"):
+        raise RuntimeError("Canon-D requested but MLP has no `.fc1`.")
+      fc1 = mlp.fc1
+      hidden_dim = getattr(fc1, "out_features", None)
+      if hidden_dim is None:
+        raise RuntimeError("Canon-D requested but could not read MLP fc1 out_features.")
+      mlp.fc1 = CanonFC1Wrapper(fc1, self._make_canon(int(hidden_dim), kernel_size, causal))
+
+    if canon_c:
+      if not hasattr(block, "mlp"):
+        raise RuntimeError("Canon-C requested but block has no `.mlp`.")
+      block.mlp = CanonInputWrapper(block.mlp, self._make_canon(int(dim), kernel_size, causal))
+
+  def forward(self, x: torch.Tensor) -> torch.Tensor:
+    return self.block(x)
+
+
+def _build_student(config: HearCanonViTConfig) -> nn.Module:
+  try:
+    import timm
+  except Exception as exc:  # noqa: BLE001
+    raise RuntimeError(f"timm is required to load HearCanonViTModel: {exc}") from exc
+
+  try:
+    model = timm.create_model(
+      config.timm_model_name,
+      img_size=tuple(config.image_size),
+      in_chans=int(config.num_channels),
+      num_classes=0,
+      global_pool="avg",
+    )
+  except Exception:
+    model = timm.create_model(
+      config.timm_model_name,
+      in_chans=int(config.num_channels),
+      num_classes=0,
+      global_pool="avg",
+    )
+
+  use_canon = bool(
+    config.canon
+    or config.canon_abcd
+    or config.canon_a
+    or config.canon_b
+    or config.canon_c
+    or config.canon_d
+  )
+  canon_a = bool(config.canon_a)
+  canon_b = bool(config.canon_b)
+  canon_c = bool(config.canon_c)
+  canon_d = bool(config.canon_d)
+  if config.canon_abcd:
+    canon_a = canon_b = canon_c = canon_d = True
+
+  if use_canon:
+    dim = getattr(model, "embed_dim", None) or getattr(model, "num_features", None)
+    if dim is None:
+      raise RuntimeError("Could not determine student embed_dim for Canon layers.")
+    if not hasattr(model, "blocks"):
+      raise RuntimeError("Student model has no `.blocks` attribute; cannot insert Canon layers.")
+
+    grid_size = None
+    expect_cls = None
+    if config.canon_2d:
+      patch_embed = getattr(model, "patch_embed", None)
+      grid_size = getattr(patch_embed, "grid_size", None) if patch_embed is not None else None
+      if grid_size is not None:
+        num_prefix = getattr(model, "num_prefix_tokens", None)
+        if num_prefix is not None:
+          try:
+            num_prefix = int(num_prefix)
+          except Exception:
+            num_prefix = None
+          if num_prefix in (0, 1):
+            expect_cls = bool(num_prefix)
+        if expect_cls is None:
+          expect_cls = getattr(model, "cls_token", None) is not None
+
+    wrapped_blocks = []
+    for block in model.blocks:
+      wrapped_blocks.append(
+        CanonBlockWrapper(
+          block,
+          int(dim),
+          kernel_size=int(config.canon_kernel),
+          canon_a=canon_a,
+          canon_b=canon_b,
+          canon_b_qkv=bool(config.canon_b_qkv),
+          canon_c=canon_c,
+          canon_d=canon_d,
+          causal=bool(config.canon_causal),
+          use_2d=bool(config.canon_2d),
+          grid_size=grid_size,
+          expect_cls=expect_cls,
+        )
+      )
+    model.blocks = nn.Sequential(*wrapped_blocks)
+
+  if bool(config.canon_no_pos_enc) and use_canon:
+    _disable_positional_embeddings(model)
+
+  return model
+
+
+def _student_features(feats: torch.Tensor) -> torch.Tensor:
+  if feats.ndim == 3:
+    return feats[:, 0, :]
+  if feats.ndim == 4:
+    return feats.mean(dim=(-2, -1))
+  return feats
+
+
+class HearCanonViTModel(PreTrainedModel):
+  """Distilled HeAR ViT-S model with Canon layers and 512-D projection head."""
+
+  config_class = HearCanonViTConfig
+  base_model_prefix = "student"
+  main_input_name = "input_values"
+
+  def __init__(self, config: HearCanonViTConfig):
+    super().__init__(config)
+    self.student = _build_student(config)
+    self.proj = nn.Linear(int(config.hidden_size), int(config.pooler_output_size))
+    self.post_init()
+
+  def preprocess_audio(self, audio: torch.Tensor) -> torch.Tensor:
+    return preprocess_audio(
+      audio,
+      clip_samples=int(self.config.num_audio_samples),
+      target_size=(int(self.config.image_size[0]), int(self.config.image_size[1])),
+    )
+
+  def embed_audio(self, audio: torch.Tensor) -> torch.Tensor:
+    with torch.inference_mode():
+      return self.forward(input_values=audio, return_dict=True).pooler_output
+
+  def _prepare_pixel_values(
+    self,
+    *,
+    input_values: Optional[torch.Tensor],
+    pixel_values: Optional[torch.Tensor],
+  ) -> torch.Tensor:
+    x = pixel_values if pixel_values is not None else input_values
+    if x is None:
+      raise ValueError("Provide `input_values` (waveform) or `pixel_values` (spectrogram).")
+
+    if not torch.is_tensor(x):
+      x = torch.tensor(x)
+
+    if x.ndim == 1:
+      x = x.unsqueeze(0)
+
+    if x.ndim == 2:
+      x = self.preprocess_audio(x)
+    elif x.ndim == 3:
+      x = x.unsqueeze(1)
+    elif x.ndim == 4:
+      if x.shape[1] != int(self.config.num_channels) and x.shape[-1] == int(self.config.num_channels):
+        x = x.permute(0, 3, 1, 2).contiguous()
+    else:
+      raise ValueError(
+        "Expected waveform [B, samples] or spectrogram [B, H, W]/[B, C, H, W]. "
+        f"Got shape {tuple(x.shape)}."
+      )
+
+    if x.ndim != 4:
+      raise ValueError(f"Expected 4D spectrogram tensor after preprocessing, got {tuple(x.shape)}.")
+
+    if x.shape[1] != int(self.config.num_channels):
+      raise ValueError(
+        f"Expected {int(self.config.num_channels)} channel(s), got {x.shape[1]}."
+      )
+
+    target_h = int(self.config.image_size[0])
+    target_w = int(self.config.image_size[1])
+    if x.shape[-2:] != (target_h, target_w):
+      x = _torch_resize_bilinear_tf_compat(x, size=(target_h, target_w))
+
+    return x.float()
+
+  def forward(
+    self,
+    input_values: Optional[torch.Tensor] = None,
+    pixel_values: Optional[torch.Tensor] = None,
+    output_hidden_states: Optional[bool] = None,
+    return_dict: Optional[bool] = None,
+    **kwargs,
+  ):
+    del kwargs
+
+    if return_dict is None:
+      return_dict = self.config.use_return_dict
+
+    x = self._prepare_pixel_values(input_values=input_values, pixel_values=pixel_values)
+    x = x.to(device=self.device)
+
+    feats = self.student.forward_features(x)
+    if isinstance(feats, (list, tuple)):
+      feats = feats[-1]
+
+    pooled_student = _student_features(feats)
+    pooler_output = self.proj(pooled_student)
+
+    if feats.ndim == 2:
+      last_hidden_state = feats.unsqueeze(1)
+    elif feats.ndim == 4:
+      b, c, h, w = feats.shape
+      last_hidden_state = feats.view(b, c, h * w).transpose(1, 2).contiguous()
+    else:
+      last_hidden_state = feats
+
+    hidden_states = (last_hidden_state,) if output_hidden_states else None
+
+    if not return_dict:
+      if output_hidden_states:
+        return (last_hidden_state, pooler_output, hidden_states)
+      return (last_hidden_state, pooler_output)
+
+    return BaseModelOutputWithPooling(
+      last_hidden_state=last_hidden_state,
+      pooler_output=pooler_output,
+      hidden_states=hidden_states,
+      attentions=None,
+    )
+
+
+__all__ = ["HearCanonViTModel", "HearCanonViTConfig", "preprocess_audio"]

preprocessor_config.json

@@ -0,0 +1,13 @@
+{
+  "clip_duration_seconds": 2.0,
+  "do_normalize": false,
+  "feature_extractor_type": "HearCanonFeatureExtractor",
+  "image_size": [
+    192,
+    128
+  ],
+  "input_channels": 1,
+  "num_audio_samples": 32000,
+  "num_mel_bins": 128,
+  "sampling_rate": 16000
+}

pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e1c0b78f1b8a42be8ed5e24f3aa807adfedbc5270d21418e7acb5c135a3c89e3
+size 89464259

requirements.txt

@@ -0,0 +1,5 @@
+torch
+transformers>=4.50.0
+timm>=1.0.0
+scipy
+soundfile

smoke_test.py

@@ -0,0 +1,40 @@
+#!/usr/bin/env python3
+"""Local smoke test for the distilled HeAR ViT-S Canon upload 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()

training_args.json

@@ -0,0 +1,69 @@
+{
+  "amp": false,
+  "batch_size": 128,
+  "canon": true,
+  "canon_2d": true,
+  "canon_a": false,
+  "canon_abcd": true,
+  "canon_b": false,
+  "canon_b_qkv": false,
+  "canon_c": false,
+  "canon_causal": false,
+  "canon_d": false,
+  "canon_kernel": 4,
+  "canon_no_pos_enc": true,
+  "canon_post": false,
+  "canon_pre": false,
+  "clip_seconds": 2.0,
+  "contrastive_temp": 0.07,
+  "data_dir": "/home/ubuntu/HeAR/HeAR/data/laion_audio_lake2",
+  "device": "cuda",
+  "gns_every": 5,
+  "gns_param_sample": 200000,
+  "grad_accum": 1,
+  "live_shard_refresh": true,
+  "log_every": 10,
+  "loss_contrastive_weight": 0.5,
+  "loss_mse_weight": 1.0,
+  "loss_relational_weight": 0.5,
+  "lr": 0.0003,
+  "lr_gns_adapt": true,
+  "lr_gns_ema_beta": 0.995,
+  "lr_gns_max_factor": 1.0,
+  "lr_gns_min_factor": 0.1,
+  "lr_gns_min_samples": 100,
+  "lr_gns_ref_batch": 0.0,
+  "lr_gns_update_every": 1000,
+  "lr_min_ratio": 0.1,
+  "lr_schedule": "none",
+  "lr_warmup_steps": 0,
+  "max_checkpoints": 20,
+  "max_steps": 200000,
+  "num_workers": 4,
+  "out": "/home/ubuntu/HeAR/HeAR/checkpoints/hear_vit_s_lake",
+  "repeat": true,
+  "resume_from": "checkpoints/hear_vit_s_lake/ckpt_013000.pt",
+  "resume_latest": false,
+  "resume_require_optim": true,
+  "sample_rate": 16000,
+  "save_every": 1000,
+  "seed": 1337,
+  "shard_refresh_sec": 20.0,
+  "shards_glob": "shard-*.tar",
+  "shuffle_shards": true,
+  "streams_glob": "stream-*",
+  "teacher_id": "google/hear-pytorch",
+  "val_batches": 5,
+  "val_defer_check_every": 10,
+  "val_defer_start_steps": 10,
+  "val_every": 250,
+  "val_fraction": 0.0,
+  "val_shards_file": "/home/ubuntu/HeAR/HeAR/checkpoints/hear_vit_s_lake/val_shards.json",
+  "val_target_clips": 10000,
+  "wandb": true,
+  "wandb_entity": null,
+  "wandb_project": "hear-distill",
+  "wandb_run_name": null,
+  "wandb_tags": null,
+  "weight_decay": 0.05
+}