README.md

---
license: mit
language:
- en
base_model:
- laion/CLIP-ViT-B-32-laion2B-s34B-b79K
- MIT/ast-finetuned-audioset-10-10-0.4593
pipeline_tag: feature-extraction
tags:
- multimodal
- cross-modal
- embeddings
- audio
- vision
- text
- contrastive-learning
- openclip
- ast
---

# 🌿 Sage Embed — v1.1

**Sage Embed** is a cross-modal embedding model that encodes **text**, **images**, and **audio** into a unified shared latent space. By aligning all three modalities, it enables semantic similarity and retrieval across any combination — query audio with text, compare images to sounds, or any mix in between.

> 🚧 **This project is in early development.** Quantitative evaluation metrics will be added in a future release. See the [Evaluation](#evaluation) section for details.

---

## Model Summary

| Property | Details |
|---|---|
| **Version** | v1.1 |
| **Modalities** | Text · Image · Audio |
| **Shared Embedding Dim** | 512 |
| **Text & Image Backbone** | [`laion/CLIP-ViT-B-32-laion2B-s34B-b79K`](https://huggingface.co/laion/CLIP-ViT-B-32-laion2B-s34B-b79K) |
| **Audio Backbone** | [`MIT/ast-finetuned-audioset-10-10-0.4593`](https://huggingface.co/MIT/ast-finetuned-audioset-10-10-0.4593) |
| **Training Dataset** | [Clotho v2](https://github.com/audio-captioning/clotho-dataset) |
| **Training Hardware** | Google Colab — NVIDIA T4 (free tier) |
| **Epochs** | 20 |
| **License** | MIT |

---

## Architecture

### Text & Image: OpenCLIP ViT-B/32
- **Model:** `hf-hub:laion/CLIP-ViT-B-32-laion2B-s34B-b79K`
- Weights are **fully frozen** during training — used purely as a fixed embedding target.
- Produces **512-dimensional** L2-normalized embeddings for both text and images.

### Audio: Audio Spectrogram Transformer (AST)
- **Model:** `MIT/ast-finetuned-audioset-10-10-0.4593`
- AST hidden size: **768**
- The model is **mostly frozen**; the **last 2 encoder layers** and the **final LayerNorm** are unfrozen for fine-tuning, allowing the audio encoder to adapt to the audio-caption alignment task without catastrophic forgetting.
- The `[CLS]` token (`last_hidden_state[:, 0]`) is extracted and fed to the projection layer.

### Audio Projection Layer

The projection layer maps AST's 768-dimensional output into CLIP's 512-dimensional shared space

---

## Evaluation

> **Metrics Coming Soon**
>
> Quantitative benchmarks will be added in a future update as the project matures. This release focuses on establishing the cross-modal architecture

---

## Usage

### Installation

```bash
pip install torch torchaudio transformers open_clip_torch huggingface_hub
```

### Step 1 — Define the Model & Load Weights

Copy this class definition exactly. The weights file is pulled directly from the Hub.

```python
import os
import torch
import torch.nn as nn
import torchaudio
import torchaudio.transforms as T
from transformers import ASTModel, AutoProcessor
from huggingface_hub import hf_hub_download
import open_clip
from PIL import Image
import numpy as np

MODEL_NAME = "hf-hub:laion/CLIP-ViT-B-32-laion2B-s34B-b79K"
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("MIT/ast-finetuned-audioset-10-10-0.4593")
        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):
        outputs = self.ast(input_values)
        # Average of first two token states for richer audio representation
        audio_features = (outputs.last_hidden_state[:, 0] + outputs.last_hidden_state[:, 1]) / 2
        return self.audio_projection(audio_features)

# --- Load weights from HuggingFace Hub ---
weights_path = hf_hub_download(
    repo_id="harryfrz/sage-embed",
    filename="sage_embed_v1.1.pt"
)

model = OpenCLIP_AST_Model().to(DEVICE)
model.load_state_dict(torch.load(weights_path, map_location=DEVICE))
model.eval()

# Processors
ast_processor = AutoProcessor.from_pretrained("MIT/ast-finetuned-audioset-10-10-0.4593")
tokenizer = open_clip.get_tokenizer(MODEL_NAME)

print("Sage Embed v1.1 loaded.")
```

---

### Step 2 — Encode Your Inputs

**Encode a single audio file** (`.wav`, `.mp3`, `.flac`):

```python
def encode_audio(filepath):
    waveform, sr = torchaudio.load(filepath)

    # Convert to mono
    if waveform.shape[0] > 1:
        waveform = torch.mean(waveform, dim=0)
    else:
        waveform = waveform.squeeze(0)

    # Resample to 16kHz
    if sr != 16000:
        waveform = T.Resample(sr, 16000)(waveform)

    # Pad or crop to 10 seconds (163840 samples @ 16kHz)
    target = 163840
    if waveform.shape[0] < target:
        waveform = torch.nn.functional.pad(waveform, (0, target - waveform.shape[0]))
    else:
        waveform = waveform[:target]

    inputs = ast_processor(waveform, sampling_rate=16000, return_tensors="pt")
    input_tensor = inputs["input_values"].to(DEVICE)

    with torch.no_grad():
        emb = model.forward_audio(input_tensor)
        emb = emb / emb.norm(dim=-1, keepdim=True)
    return emb  # shape: [1, 512]


audio_emb = encode_audio("your_audio.wav")
```

**Encode a single image file** (`.jpg`, `.png`, `.webp`):

```python
def encode_image(filepath):
    image = Image.open(filepath).convert("RGB")
    image_input = model.image_preprocess(image).unsqueeze(0).to(DEVICE)

    with torch.no_grad():
        emb = model.clip_model.encode_image(image_input)
        emb = emb / emb.norm(dim=-1, keepdim=True)
    return emb  # shape: [1, 512]


image_emb = encode_image("your_image.jpg")
```

**Encode text:**

```python
def encode_text(texts):
    # texts: a string or list of strings
    if isinstance(texts, str):
        texts = [texts]
    tokens = tokenizer(texts).to(DEVICE)

    with torch.no_grad():
        emb = model.clip_model.encode_text(tokens)
        emb = emb / emb.norm(dim=-1, keepdim=True)
    return emb  # shape: [N, 512]


text_emb = encode_text("a dog barking in the park")
```

---

### Step 3 — Cross-Modal Similarity Test

All embeddings are **L2-normalized**, so dot product equals cosine similarity directly.

```python
# --- Audio ↔ Text ---
audio_emb = encode_audio("dog_bark.wav")
text_embs = encode_text(["a dog barking", "rain on a rooftop", "busy city traffic"])

sim = (audio_emb @ text_embs.T)[0]  # shape: [3]
labels = ["a dog barking", "rain on a rooftop", "busy city traffic"]

print("Audio ↔ Text similarity:")
for label, score in zip(labels, sim.tolist()):
    print(f"  '{label}': {score:.4f}")
# Expected: highest score for "a dog barking"


# --- Image ↔ Text ---
image_emb = encode_image("sunset.jpg")
text_embs = encode_text(["a sunset over the ocean", "a crowded street", "birds chirping"])

sim = (image_emb @ text_embs.T)[0]
labels = ["a sunset over the ocean", "a crowded street", "birds chirping"]

print("\nImage ↔ Text similarity:")
for label, score in zip(labels, sim.tolist()):
    print(f"  '{label}': {score:.4f}")


# --- Audio ↔ Image (cross-modal) ---
sim_cross = (audio_emb @ image_emb.T)[0]
print(f"\n Audio ↔ Image similarity: {sim_cross.item():.4f}")
```

---

## Limitations

- **Early-stage model:** Sage Embed v1.1 is the first public release. The projection layer was trained on Clotho only; performance on out-of-distribution audio domains may vary.
- **CLIP is frozen:** The text and image encoders are not jointly fine-tuned, so alignment quality is bounded by how well the audio projection adapts to CLIP's fixed representation space.
- **Training scale:** Trained on a free-tier T4 GPU with a relatively small dataset. Larger-scale training is expected to improve retrieval quality significantly in future versions.
- **Audio length:** Optimized for clips up to ~10 seconds. Longer clips are truncated from the start.

---

## Acknowledgements

- [OpenCLIP](https://github.com/mlfoundations/open_clip) — LAION CLIP ViT-B/32 backbone
- [Audio Spectrogram Transformer](https://huggingface.co/MIT/ast-finetuned-audioset-10-10-0.4593) — MIT AST backbone
- [Clotho Dataset](https://github.com/audio-captioning/clotho-dataset) — Training data