whisper-clap-version-0.1

import torch
import torch.nn as nn
import torch.nn.functional as F
import torchaudio
from transformers import WhisperModel, WhisperFeatureExtractor, AutoTokenizer, AutoModel
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file

# --- Model Definition ---
class WhisperClapModel(nn.Module):
    def __init__(self, whisper_name="openai/whisper-small"):
        super().__init__()
        self.audio_encoder = WhisperModel.from_pretrained(whisper_name).encoder
        self.projector = nn.Sequential(
            nn.Linear(768, 2048),
            nn.GELU(),
            nn.Linear(2048, 768)
        )

    def forward(self, input_features):
        # Extract features and mean pool
        outputs = self.audio_encoder(input_features)
        rep = outputs.last_hidden_state.mean(dim=1) 
        # Project to GTE space and normalize
        emb = self.projector(rep)
        return F.normalize(emb, p=2, dim=1)

# --- Setup ---
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
WHISPER_NAME = "openai/whisper-small"
TEXT_MODEL_NAME = "Alibaba-NLP/gte-base-en-v1.5"

# 1. Load Model
model = WhisperClapModel(WHISPER_NAME).to(DEVICE)

# Download and load weights
checkpoint_path = hf_hub_download(repo_id="laion/whisper-clap-version-0.1", filename="model.safetensors")
state_dict = load_file(checkpoint_path)

# Clean keys (remove 'model.' prefix if present from training)
clean_state = {k.replace("model.", ""): v for k, v in state_dict.items()}
model.load_state_dict(clean_state, strict=False)
model.eval()

# 2. Load Processors
feature_extractor = WhisperFeatureExtractor.from_pretrained(WHISPER_NAME)
tokenizer = AutoTokenizer.from_pretrained(TEXT_MODEL_NAME)
text_model = AutoModel.from_pretrained(TEXT_MODEL_NAME, trust_remote_code=True).to(DEVICE).eval()

# 3. Helper Functions
def encode_audio(audio_path):
    wav, sr = torchaudio.load(audio_path)
    # Resample to 16k
    if sr != 16000:
        wav = torchaudio.functional.resample(wav, sr, 16000)
    # Mono
    if wav.shape[0] > 1:
        wav = torch.mean(wav, dim=0, keepdim=True)
    # Pad/Trim to 30s
    target_len = 16000 * 30
    if wav.shape[1] > target_len:
        wav = wav[:, :target_len]
    elif wav.shape[1] < target_len:
        wav = F.pad(wav, (0, target_len - wav.shape[1]))
        
    inputs = feature_extractor(wav.squeeze().numpy(), sampling_rate=16000, return_tensors="pt")
    with torch.no_grad():
        return model(inputs.input_features.to(DEVICE))

def encode_text(text):
    inputs = tokenizer(text, max_length=512, padding=True, truncation=True, return_tensors="pt").to(DEVICE)
    with torch.no_grad():
        outputs = text_model(**inputs)
        # GTE pooling
        mask = inputs.attention_mask.unsqueeze(-1).expand(outputs.last_hidden_state.size()).float()
        sum_embeddings = torch.sum(outputs.last_hidden_state * mask, 1)
        sum_mask = torch.clamp(mask.sum(1), min=1e-9)
        return F.normalize(sum_embeddings / sum_mask, p=2, dim=1)

# --- Run ---
# audio_emb = encode_audio("my_audio.wav")
# text_emb = encode_text("A calm, soothing voice explaining a concept.")
# similarity = torch.matmul(audio_emb, text_emb.t()).item()
# print(f"Similarity: {similarity:.4f}")

import os
import json
import glob
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchaudio
from torch.utils.data import Dataset, DataLoader
from transformers import WhisperModel, WhisperFeatureExtractor, AutoTokenizer, AutoModel

# --- Configuration ---
DATA_FOLDER = "./my_training_data"  # Path to your data
BATCH_SIZE = 16
LEARNING_RATE = 1e-4
EPOCHS = 3
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

# --- 1. Dataset Loader ---
class AudioTextDataset(Dataset):
    def __init__(self, folder_path):
        self.samples = []
        # Find all audio files (mp3 and wav)
        audio_files = glob.glob(os.path.join(folder_path, "*.[mM][pP]3")) + \
                      glob.glob(os.path.join(folder_path, "*.[wW][aA][vV]"))
        
        for audio_path in audio_files:
            # Construct expected JSON path (same basename)
            base_path = os.path.splitext(audio_path)[0]
            json_path = base_path + ".json"
            
            if os.path.exists(json_path):
                try:
                    with open(json_path, 'r', encoding='utf-8') as f:
                        data = json.load(f)
                        # Ensure your JSON has a 'caption' key, or adjust this line
                        caption = data.get("caption") 
                        if caption:
                            self.samples.append((audio_path, caption))
                except Exception as e:
                    print(f"Error loading {json_path}: {e}")

    def __len__(self):
        return len(self.samples)

    def __getitem__(self, idx):
        audio_path, text = self.samples[idx]
        
        # Load and process Audio
        wav, sr = torchaudio.load(audio_path)
        if sr != 16000:
            wav = torchaudio.functional.resample(wav, sr, 16000)
        # Mix to mono if necessary
        if wav.shape[0] > 1:
            wav = torch.mean(wav, dim=0, keepdim=True)
            
        # Random crop or pad to 30s (simple strategy for training)
        target_len = 16000 * 30
        if wav.shape[1] > target_len:
            # Random crop
            start = torch.randint(0, wav.shape[1] - target_len, (1,)).item()
            wav = wav[:, start:start+target_len]
        elif wav.shape[1] < target_len:
            wav = F.pad(wav, (0, target_len - wav.shape[1]))

        return {"audio": wav.squeeze(), "text": text}

# --- 2. Training Model Wrapper ---
class WhisperClapTrainer(nn.Module):
    def __init__(self):
        super().__init__()
        # Audio Encoder (Trainable)
        self.whisper = WhisperModel.from_pretrained("openai/whisper-small")
        self.audio_encoder = self.whisper.encoder
        
        # Projection Head (Trainable)
        self.projector = nn.Sequential(
            nn.Linear(768, 2048),
            nn.GELU(),
            nn.Linear(2048, 768)
        )
        
        # Text Encoder (Frozen - we want to align audio to the fixed GTE space)
        self.text_model = AutoModel.from_pretrained("Alibaba-NLP/gte-base-en-v1.5", trust_remote_code=True)
        for p in self.text_model.parameters():
            p.requires_grad = False

    def forward(self, input_features, input_ids, attention_mask):
        # 1. Encode Audio
        audio_out = self.audio_encoder(input_features)
        audio_rep = audio_out.last_hidden_state.mean(dim=1)
        audio_emb = F.normalize(self.projector(audio_rep), p=2, dim=1)

        # 2. Encode Text (with Gradient disabled)
        with torch.no_grad():
            text_out = self.text_model(input_ids=input_ids, attention_mask=attention_mask)
            # GTE specific pooling logic
            mask_expanded = attention_mask.unsqueeze(-1).expand(text_out.last_hidden_state.size()).float()
            sum_embeddings = torch.sum(text_out.last_hidden_state * mask_expanded, 1)
            sum_mask = torch.clamp(mask_expanded.sum(1), min=1e-9)
            text_emb = F.normalize(sum_embeddings / sum_mask, p=2, dim=1)

        # 3. Compute Loss (MSE or Cosine Embedding Loss)
        # We maximize cosine similarity (minimize 1 - cos_sim)
        cos_sim = F.cosine_similarity(audio_emb, text_emb)
        loss = 1.0 - cos_sim.mean()
        
        return loss

# --- 3. Main Training Loop ---
def train():
    # Tools
    feature_extractor = WhisperFeatureExtractor.from_pretrained("openai/whisper-small")
    tokenizer = AutoTokenizer.from_pretrained("Alibaba-NLP/gte-base-en-v1.5")
    
    # Data
    dataset = AudioTextDataset(DATA_FOLDER)
    dataloader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=4)
    
    # Model
    model = WhisperClapTrainer().to(DEVICE)
    optimizer = torch.optim.AdamW(filter(lambda p: p.requires_grad, model.parameters()), lr=LEARNING_RATE)

    model.train()
    print(f"Starting training on {len(dataset)} pairs...")

    for epoch in range(EPOCHS):
        total_loss = 0
        for batch in dataloader:
            optimizer.zero_grad()
            
            # Prepare Inputs
            audio_inputs = feature_extractor(batch["audio"].numpy(), sampling_rate=16000, return_tensors="pt")
            text_inputs = tokenizer(batch["text"], padding=True, truncation=True, max_length=512, return_tensors="pt")
            
            # Move to device
            input_features = audio_inputs.input_features.to(DEVICE)
            input_ids = text_inputs.input_ids.to(DEVICE)
            attn_mask = text_inputs.attention_mask.to(DEVICE)

            # Forward & Backward
            loss = model(input_features, input_ids, attn_mask)
            loss.backward()
            optimizer.step()
            
            total_loss += loss.item()
            
        print(f"Epoch {epoch+1}/{EPOCHS} - Avg Loss: {total_loss / len(dataloader):.4f}")

    # Save final model
    torch.save(model.state_dict(), "whisper_clap_finetuned.pt")
    print("Training complete. Model saved.")

if __name__ == "__main__":
    train()