Add model card

README.md

@@ -1,3 +1,312 @@
 ---
 license: cc-by-4.0
+library_name: pytorch
+pipeline_tag: feature-extraction
+tags:
+  - audio-text-retrieval
+  - contrastive-learning
+  - clip
+  - openclip
+  - audio-spectrogram-transformer
+  - cross-modal-alignment
+  - multimodal-embeddings
+base_model: MIT/ast-finetuned-audioset-10-10-0.4593
 ---
+
+# Edgebind
+
+Audio-into-CLIP alignment for text–image–audio retrieval. Edgebind maps audio into the
+existing embedding space of a frozen CLIP model instead of training a joint encoder
+from scratch, so audio, text and images become directly comparable by cosine
+similarity in one 512-d space.
+
+Research artifact for **"Edgebind: Towards Edge-Compatible Audio-into-CLIP Alignment
+for Text-Image-Audio Retrieval"** (CAISc 2026).
+
+> **Note on the `datasets` metadata field.** This card intentionally omits it. Clotho
+> v2.1 has no canonical publisher-owned dataset repository on the Hub — a search
+> returns only third-party mirrors and the unrelated ClothoAQA task. The training code
+> downloads Clotho directly from Zenodo, so tagging a mirror would misstate
+> provenance. The authoritative source is linked in prose under
+> [Training data](#training-data).
+
+## Model description
+
+| Component | Detail |
+|---|---|
+| Text & image encoders | OpenCLIP ViT-B/32, `hf-hub:laion/CLIP-ViT-B-32-laion2B-s34B-b79K` — **fully frozen**, never updated |
+| Audio encoder | Audio Spectrogram Transformer, [`MIT/ast-finetuned-audioset-10-10-0.4593`](https://huggingface.co/MIT/ast-finetuned-audioset-10-10-0.4593) |
+| AST frozen | Patch embeddings, position embeddings, encoder layers 0–8 |
+| AST trained | Encoder layers 9, 10, 11 and the final layernorm |
+| Pooling | Mean of the first two AST output tokens (CLS and distillation) |
+| Projection head | `Linear(768, 1024) → LayerNorm(1024) → ReLU → Dropout(0.3) → Linear(1024, 512)` |
+| Temperature | Learnable `logit_scale`, initialized to `log(1 / 0.07)` |
+| Output | 512-d, L2-normalized, shared with CLIP text and image embeddings |
+| Objective | Symmetric InfoNCE (cross-entropy over in-batch negatives in both directions, averaged) |
+| Training data | Clotho v2.1, development split |
+
+Only the AST upper layers, the projection head and the temperature receive gradients.
+Because the CLIP towers are untouched, text and image embeddings produced by this model
+are identical to stock OpenCLIP ViT-B/32 — the alignment is carried entirely by the
+audio branch.
+
+The pooling choice is worth stating explicitly, since it is easy to assume otherwise:
+audio features are the **average of tokens 0 and 1**, not the CLS token alone. Using
+CLS alone is a separate ablation in the training notebook.
+
+## Intended use and limitations
+
+Intended use: text-to-audio and audio-to-text retrieval over a local corpus you
+control — embed a set of audio clips once, then rank them against free-text queries by
+cosine similarity. Because the CLIP image tower is frozen and shared, images can be
+embedded into the same space and searched with the same text queries.
+
+**This is a research prototype, not a production model.** Specifically:
+
+- **Single dataset.** Trained and validated only on Clotho v2.1, which is small
+  (a few thousand clips) and skewed toward everyday environmental and ambient sound.
+  Behaviour on speech, music, or domain-specific audio is uncharacterized.
+- **Single training run, no seed variance.** One run of 20 epochs. The training code
+  sets no random seed, and caption sampling, crop offsets and SpecAugment masks are all
+  stochastic, so run-to-run variance has not been measured. Treat any single reported
+  number as one sample, not a mean.
+- **No CPU, quantized, or on-device benchmarking.** The model was trained and run under
+  CUDA mixed precision on a single NVIDIA T4. Latency, memory, and accuracy under CPU
+  inference, quantization, distillation, or mobile/embedded runtimes have **not** been
+  evaluated. The released checkpoint is unquantized fp32. Despite "edge-compatible" in
+  the paper title — which refers to the design motivation for reusing a frozen
+  ViT-B/32 backbone — nothing here establishes that this model is edge-ready or
+  deployable on constrained hardware. Do not treat it as such.
+- **Fixed 10.24 s window.** Audio is cropped or zero-padded to exactly 163,840 samples
+  at 16 kHz. Longer recordings are truncated, not chunked; content outside the window
+  is invisible to the model.
+- **Inherited bias.** The text and image behaviour is entirely that of
+  LAION-2B-trained OpenCLIP ViT-B/32 and carries its biases unchanged.
+
+Out of scope: audio captioning or generation (there is no decoder), speaker or speech
+recognition, and any safety-, surveillance-, or identity-related classification.
+
+## How to use
+
+The checkpoint is a `state_dict` for the composite module defined in the training
+notebook, so you must reconstruct that module — CLIP submodule included — before
+loading. The snippet below mirrors the notebook's own model definition and its
+evaluation-time preprocessing path.
+
+```python
+import numpy as np
+import torch
+import torch.nn as nn
+import torchaudio
+import torchaudio.transforms as T
+from huggingface_hub import hf_hub_download
+from transformers import ASTModel, AutoProcessor
+import open_clip
+
+MODEL_NAME = "hf-hub:laion/CLIP-ViT-B-32-laion2B-s34B-b79K"
+AST_NAME = "MIT/ast-finetuned-audioset-10-10-0.4593"
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+
+class OpenCLIP_AST_Model(nn.Module):
+    def __init__(self, embedding_dim=512):
+        super().__init__()
+        self.ast = ASTModel.from_pretrained(AST_NAME)
+        self.clip_model, _, self.image_preprocess = open_clip.create_model_and_transforms(
+            MODEL_NAME, pretrained=None
+        )
+        self.audio_projection = nn.Sequential(
+            nn.Linear(self.ast.config.hidden_size, 1024),
+            nn.LayerNorm(1024),
+            nn.ReLU(),
+            nn.Dropout(0.3),
+            nn.Linear(1024, embedding_dim),
+        )
+        self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
+
+    def forward_audio(self, input_values):
+        out = self.ast(input_values)
+        # Mean of CLS + distillation tokens, matching training.
+        feats = (out.last_hidden_state[:, 0] + out.last_hidden_state[:, 1]) / 2
+        return self.audio_projection(feats)
+
+
+weights = hf_hub_download("harryfrz/edgebind", "edgebind_v1.1")
+
+model = OpenCLIP_AST_Model().to(DEVICE)
+# The notebook loads with a plain torch.load. On torch >= 2.6 the weights_only=True
+# default is appropriate for a pure tensor state_dict; pass weights_only=False if
+# your torch version raises on it.
+model.load_state_dict(torch.load(weights, map_location=DEVICE))
+model.eval()
+
+ast_processor = AutoProcessor.from_pretrained(AST_NAME)
+tokenizer = open_clip.get_tokenizer(MODEL_NAME)
+
+SAMPLE_RATE = 16000
+TARGET_SAMPLES = 163840  # 10.24 s at 16 kHz
+
+
+def embed_audio(path):
+    waveform, sr = torchaudio.load(path)
+    waveform = waveform.mean(dim=0) if waveform.shape[0] > 1 else waveform.squeeze(0)
+    if sr != SAMPLE_RATE:
+        waveform = T.Resample(sr, SAMPLE_RATE)(waveform)
+
+    n = waveform.shape[0]
+    if n < TARGET_SAMPLES:
+        waveform = torch.nn.functional.pad(waveform, (0, TARGET_SAMPLES - n))
+    elif n > TARGET_SAMPLES:
+        start = (n - TARGET_SAMPLES) // 2  # centre crop, as at evaluation time
+        waveform = waveform[start:start + TARGET_SAMPLES]
+
+    inputs = ast_processor(waveform, sampling_rate=SAMPLE_RATE, return_tensors="pt")
+    with torch.no_grad():
+        emb = model.forward_audio(inputs["input_values"].to(DEVICE))
+    return emb / emb.norm(dim=-1, keepdim=True)
+
+
+def embed_text(prompts):
+    with torch.no_grad():
+        emb = model.clip_model.encode_text(tokenizer(prompts).to(DEVICE))
+    return emb / emb.norm(dim=-1, keepdim=True)
+
+
+def embed_image(pil_image):
+    with torch.no_grad():
+        x = model.image_preprocess(pil_image).unsqueeze(0).to(DEVICE)
+        emb = model.clip_model.encode_image(x)
+    return emb / emb.norm(dim=-1, keepdim=True)
+
+
+# --- retrieval over a local corpus ---
+corpus = ["clip_a.wav", "clip_b.wav", "clip_c.wav"]
+index = torch.cat([embed_audio(p) for p in corpus], dim=0)   # [N, 512]
+
+query = embed_text(["waves hitting the shore"])              # [1, 512]
+scores = (query @ index.T)[0]                                # cosine similarity
+
+for rank in scores.argsort(descending=True):
+    print(f"{corpus[rank]}  {scores[rank].item():.4f}")
+```
+
+`open_clip.create_model_and_transforms(MODEL_NAME, pretrained=None)` still downloads
+pretrained weights: the `hf-hub:` prefix resolves the checkpoint from the Hub, and
+`pretrained=None` only means no additional named tag is applied. Those weights are
+then overwritten by the `clip_model.*` entries in the state dict.
+
+### What the checkpoint file contains
+
+`edgebind_v1.1` (955 MB, no file extension). The training code saves with
+`torch.save(model.state_dict(), ...)` and reloads with a strict
+`model.load_state_dict(torch.load(path, map_location=device))`. That means:
+
+- It is a **raw `state_dict`**, not a training checkpoint — no optimizer state, no
+  epoch counter, no scheduler or scaler state, no embedded config or metrics.
+- It stores the **complete** module, not only the trained tensors. Key prefixes are
+  `ast.*` (the whole AST, frozen layers 0–8 included), `clip_model.*` (the entire
+  frozen OpenCLIP ViT-B/32 image *and* text towers), `audio_projection.*`, plus the
+  scalar `logit_scale`.
+- That is roughly 239 M parameters in fp32 (AST ≈86 M + CLIP ≈151 M + projection
+  ≈1.3 M), consistent with the 955 MB file size.
+- Loading is strict, so the module must be rebuilt exactly — including the CLIP
+  submodule, which is why the snippet above instantiates it.
+
+**This description is derived from the notebook's save/load code and the file size, not
+from opening the file.** The tensor keys and dtypes have not been enumerated directly.
+If you need that confirmed, load it and inspect `.keys()`.
+
+## Evaluation
+
+**No results table is published in this card.** The training notebook was committed
+with all cell outputs stripped, and it contains no code that computes R@1/R@5/R@10,
+MedR, MeanR or mAP@10, no CLAP baseline, and no frozen-AST ablation. Transcribing
+metrics from any other source would not be verifiable against this repository, so none
+are reproduced here.
+
+For the reported text-to-audio and audio-to-text results on the Clotho v2.1 evaluation
+split, and the baseline and ablation comparisons, see the CAISc 2026 paper.
+
+The notebook does include a qualitative check — top-3 retrieval for five hardcoded
+prompts — and one ablation, CLS-only pooling, described below.
+
+## Training details
+
+### Training data
+
+Clotho v2.1, **development** split, obtained from Zenodo record
+[4783391](https://zenodo.org/records/4783391). The training code downloads
+`clotho_captions_development.csv` and `clotho_audio_development.7z` from that record
+directly. Each audio file carries five human-written captions; one is sampled at random
+per example per epoch during training.
+
+### Preprocessing
+
+- Downmix to mono, resample to 16 kHz.
+- Crop or zero-pad to 163,840 samples (10.24 s): **random** crop during training,
+  **centre** crop at evaluation.
+- Log-mel features via `AutoProcessor` for `MIT/ast-finetuned-audioset-10-10-0.4593`.
+- SpecAugment, training only: `FrequencyMasking(freq_mask_param=24)` and
+  `TimeMasking(time_mask_param=40)`.
+
+### Hyperparameters
+
+| Setting | Value |
+|---|---|
+| Optimizer | AdamW, `weight_decay=0.05` |
+| Epochs | 20 |
+| Batch size | 64 (gradient accumulation steps = 1, so effective batch = 64) |
+| LR — AST layers 9–11 | 5e-6 |
+| LR — AST final layernorm | 5e-6 |
+| LR — projection head | 2e-4 |
+| LR — logit scale | 2e-4 |
+| Schedule | `CosineAnnealingLR`, `eta_min=1e-6`, stepped per optimizer step |
+| Precision | CUDA mixed precision (`torch.amp.autocast` + `GradScaler`) |
+| Loader | `drop_last=True`, `num_workers=4` |
+
+### Compute
+
+Single NVIDIA T4. The training cell's recorded execution window spans **50 min 25 s**
+for all 20 epochs plus audio-embedding export. That figure comes from notebook cell
+execution timestamps, not from a printed training log.
+
+### Ablation included in the code
+
+**CLS-only pooling** — identical in every other respect, but uses
+`last_hidden_state[:, 0]` instead of the mean of tokens 0 and 1. It trains to a
+separate checkpoint. No metrics for it are present in the notebook.
+
+## A note on naming
+
+Some artifacts and in-code comments use the internal name **"Sage-Embed"** (for
+example, `# Verified Sage-Embed v1.1 train/freeze configuration`). This refers to the
+same model as Edgebind. The `v1.1` suffix on the checkpoint filename corresponds to
+that internal versioning.
+
+## Links
+
+- **Code:** https://github.com/harryfrzz/edgebind
+- **Paper:** "Edgebind: Towards Edge-Compatible Audio-into-CLIP Alignment for
+  Text-Image-Audio Retrieval", CAISc 2026
+- **Training data:** Clotho v2.1 — https://zenodo.org/records/4783391
+- **Base audio model:** https://huggingface.co/MIT/ast-finetuned-audioset-10-10-0.4593
+- **Base CLIP model:** https://huggingface.co/laion/CLIP-ViT-B-32-laion2B-s34B-b79K
+
+## Citation
+
+```bibtex
+@inproceedings{edgebind2026,
+  title     = {Edgebind: Towards Edge-Compatible Audio-into-CLIP Alignment for
+               Text-Image-Audio Retrieval},
+  author    = {Harikrishna C},
+  booktitle = {CAISc},
+  year      = {2026}
+}
+```
+
+The repository records no full author list, DOI, or page numbers; the author field
+above is taken from commit metadata and should be completed before use.
+
+## License
+
+CC BY 4.0.