beabucayan/lingkodai-audio-lid-head

LingkodAI Audio LID Head

Identifies whether a spoken clip is Cebuano, Filipino, or English, so the LingkodAI pipeline can route it to the appropriate Whisper STT model.

How it works

A small classifier head on top of the frozen openai/whisper-medium encoder. The encoder's output is averaged into 50 chunks (0.6 s each), and the head uses attention pooling plus a 3-layer MLP to predict the language.

Classes: 0 = ceb, 1 = fil, 2 = eng

Files

• deeper_50chunks.pt: weights for the classifier head only. The whisper-medium encoder is loaded separately.

Usage

import numpy as np, torch, torch.nn as nn, librosa
from transformers import WhisperFeatureExtractor, WhisperModel
from huggingface_hub import hf_hub_download

LABELS = ("ceb", "fil", "eng")
device = "cuda" if torch.cuda.is_available() else "cpu"

def build_head():
    return nn.ModuleDict({
        "attn_fc1": nn.Linear(1024, 128),
        "attn_fc2": nn.Linear(128, 1),
        "clf_fc1":  nn.Linear(2048, 256),
        "clf_fc2":  nn.Linear(256, 128),
        "clf_fc3":  nn.Linear(128, 3),
    })

def head_forward(head, chunks, num_valid):
    scores = head["attn_fc2"](torch.tanh(head["attn_fc1"](chunks))).squeeze(-1)
    idx = torch.arange(chunks.shape[1], device=chunks.device).unsqueeze(0)
    scores = scores.masked_fill(idx >= num_valid.unsqueeze(1), float("-inf"))
    w = torch.softmax(scores, dim=1).unsqueeze(-1)
    mean = (chunks * w).sum(1)
    std = torch.sqrt((((chunks - mean.unsqueeze(1)) ** 2) * w).sum(1).clamp(min=1e-8))
    h = torch.relu(head["clf_fc1"](torch.cat([mean, std], 1)))
    h = torch.relu(head["clf_fc2"](h))
    return head["clf_fc3"](h)

# Load the frozen encoder (fp32) and the head
fe = WhisperFeatureExtractor.from_pretrained("openai/whisper-medium")
encoder = WhisperModel.from_pretrained("openai/whisper-medium", torch_dtype=torch.float32).encoder.to(device).eval()
head = build_head()
path = hf_hub_download("lingkodai/lingkodai-audio-lid-head", "deeper_50chunks.pt")
head.load_state_dict(torch.load(path, map_location=device, weights_only=True))
head = head.to(device).eval()

# Predict on a 16 kHz mono clip
audio, _ = librosa.load("clip.wav", sr=16000)
feats = fe(audio, sampling_rate=16000, return_tensors="pt").input_features.to(device)
with torch.no_grad():
    hidden = encoder(feats).last_hidden_state
    chunks = hidden.view(1, 50, 30, 1024).mean(dim=2)         # 50 chunks of 30 frames
    real_frames = min(int(np.round(len(audio) / 16000 * 50)), 1500)
    num_valid = torch.tensor([int(np.clip(np.ceil(real_frames / 30), 1, 50))], device=device)
    probs = torch.softmax(head_forward(head, chunks, num_valid), dim=1)[0]

print(LABELS[int(probs.argmax())], probs.tolist())

The chunking (50 × 30 frames), valid-chunk masking, and fp32 encoder must match these settings, since the head was trained under them.