Switch to student-only 384D embeddings and update model card

README.md

@@ -13,37 +13,50 @@ tags:
 - canon
 ---
 
-# Distilled HeAR ViT-S Canon (PyTorch)
+# Distilled HeAR ViT-S Canon model card
 
-This repository package contains a Hugging Face-compatible export of our distilled HeAR student:
+**Model documentation:** HeAR (Google Health Acoustic Representations)
 
-- 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]`
+## Model information
 
-This folder is ready for upload to Hugging Face Hub as-is.
+This package contains a distilled HeAR student model implemented in PyTorch with a ViT-S backbone and Canon layers.
+
+### Description
+
+The model is built for health-acoustic embedding extraction from short audio clips.
+
+- Backbone: ViT-S (`vit_small_patch16_224`)
+- Input: single-channel mel+PCEN spectrograms (`[B, 1, 192, 128]`) generated from 2-second audio clips at 16 kHz
+- Canon setup: A/B/C/D enabled, 2D Canon, kernel size 4, positional encodings disabled
+- Output embedding: `pooler_output` with shape `[B, 384]`
 
 ## Files in this package
 
-- `config.json`: model config + `auto_map` for custom loading
+- `config.json`: model config and `auto_map`
 - `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
+- `modeling_hear_canon.py`: custom `PreTrainedModel` with integrated audio preprocessing
+- `pytorch_model.bin`: distilled student 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
+- `model_shapes.json`: structure and tensor shape inventory
+- `training_args.json`: training/checkpoint args captured from the source checkpoint
+- `.gitattributes`: git/LFS attributes for model artifacts
 - `smoke_test.py`: local verification script
 
-## Quick start (local folder)
+## How to use
+
+Install dependencies:
 
 ```bash
 pip install -U "transformers>=4.50.0" timm torch scipy soundfile
+```
+
+Run local smoke test:
+
+```bash
 python3 trained_model_hf_upload/smoke_test.py
 ```
 
-## Inference from raw audio waveform
+### Inference from raw audio waveform
 
 ```python
 import torch
@@ -62,10 +75,10 @@ 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])
+print(embeddings.shape)  # torch.Size([4, 384])
 ```
 
-## Inference from `.wav` file
+### Inference from `.wav` file
 
 ```python
 import torch
@@ -92,10 +105,10 @@ waveform = load_wav_mono_16k("example.wav")
 with torch.inference_mode():
     embedding = model.embed_audio(waveform)
 
-print(embedding.shape)  # torch.Size([1, 512])
+print(embedding.shape)  # torch.Size([1, 384])
 ```
 
-## Inference from preprocessed spectrograms
+### Inference from preprocessed spectrograms
 
 ```python
 import torch
@@ -110,36 +123,15 @@ 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 .
+print(spectrogram.shape)        # torch.Size([2, 1, 192, 128])
+print(out.pooler_output.shape)  # torch.Size([2, 384])
 ```
 
-## Architecture summary
+## Model architecture overview
 
-- Student parameters: `22,140,288`
-- Projection head parameters: `197,120`
-- Total parameters: `22,337,408`
-- Student hidden size: `384`
-- Pooler output size: `512`
+- Student model parameters: `22,140,288`
+- Embedding dimension: `384`
+- Input shape: `[B, 1, 192, 128]`
+- Output shape: `[B, 384]`
 
-Detailed tensor shapes are listed in `model_shapes.json`.
+Detailed tensor shapes are provided in `model_shapes.json`.

config.json

@@ -33,8 +33,8 @@
   "num_channels": 1,
   "num_hidden_layers": 12,
   "patch_size": 16,
-  "pooled_dim": 512,
-  "pooler_output_size": 512,
+  "pooled_dim": 384,
+  "pooler_output_size": 384,
   "sample_rate": 16000,
   "timm_model_name": "vit_small_patch16_224",
   "torch_dtype": "float32",

configuration_hear_canon.py

@@ -22,8 +22,8 @@ class HearCanonViTConfig(PretrainedConfig):
     num_hidden_layers: int = 12,
     num_attention_heads: int = 6,
     intermediate_size: int = 1536,
-    pooled_dim: int = 512,
-    pooler_output_size: int = 512,
+    pooled_dim: int = 384,
+    pooler_output_size: int = 384,
     hidden_act: str = "gelu",
     layer_norm_eps: float = 1e-6,
     sample_rate: int = 16000,

model_shapes.json

@@ -1,5 +1,6 @@
 {
   "checkpoint_source": "/home/matt/Documents/HeAR/trained_model.pt",
+  "embedding_dimension": 384,
   "exported_weights_file": "pytorch_model.bin",
   "input_spectrogram_shape": [
     "B",
@@ -11,19 +12,6 @@
     "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": [
@@ -1349,5 +1337,5 @@
     ]
   },
   "student_state_dict_size": 342,
-  "total_parameters": 22337408
+  "total_parameters": 22140288
 }

modeling_hear_canon.py

@@ -591,7 +591,7 @@ def _student_features(feats: torch.Tensor) -> torch.Tensor:
 
 
 class HearCanonViTModel(PreTrainedModel):
-  """Distilled HeAR ViT-S model with Canon layers and 512-D projection head."""
+  """Distilled HeAR ViT-S model with Canon layers and 384-D student embeddings."""
 
   config_class = HearCanonViTConfig
   base_model_prefix = "student"
@@ -600,7 +600,6 @@ class HearCanonViTModel(PreTrainedModel):
   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:
@@ -678,8 +677,7 @@ class HearCanonViTModel(PreTrainedModel):
     if isinstance(feats, (list, tuple)):
       feats = feats[-1]
 
-    pooled_student = _student_features(feats)
-    pooler_output = self.proj(pooled_student)
+    pooler_output = _student_features(feats)
 
     if feats.ndim == 2:
       last_hidden_state = feats.unsqueeze(1)

pytorch_model.bin

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

smoke_test.py

@@ -1,5 +1,5 @@
 #!/usr/bin/env python3
-"""Local smoke test for the distilled HeAR ViT-S Canon upload package."""
+"""Local smoke test for the distilled HeAR ViT-S Canon model package."""
 
 from __future__ import annotations