jaxcld 0.2.4

CLD: language detection heads for ASR models (Whisper / MMS)

CLD Convex Low-resource Accent-Robust Language Detection in Speech Recognition A lightweight language-detection module for multilingual ASR, optimized via ADMM in JAX.

---

This repository provides the official implementation of CLD, a lightweight language-detection module for multilingual ASR. This codebase contains our pip-installable Python package (jaxcld/) including our training/benchmark scripts implemented in JAX and optimized via ADMM for high performance in low-resource settings. Simply, the package attaches a small language detection head (Convex NN / small NN / linear SVM) to ASR encoder representations, and use it to select the language token (Whisper) or adapter (MMS) before decoding.

Highlights

• High Accuracy: Excels in binary and multiclass language detection.
• Low-Resource Robustness: Effective with limited data.
• Efficient: 13x training speedup from traditional NNs due to ADMM optimization and JAX.

Requirements

The package is published on PyPI as jaxcld. JAX (CPU) is included by default. For GPU support, install the matching extra:

# CPU
pip install jaxcld

# GPU — CUDA 12 (most modern systems)
pip install "jaxcld[cuda12]"

# GPU — CUDA 11
pip install "jaxcld[cuda11]"

If you've cloned this repo, you can instead install from source:

• Package-only install (inference usage):

pip install -e .                  # CPU
pip install -e ".[cuda12]"        # GPU — CUDA 12

• Full training/benchmark environment (recommended if you run the scripts in this repo):

pip install -e ".[train]"         # CPU
pip install -e ".[train,cuda12]"  # GPU — CUDA 12

If you prefer installing from the pinned dependency list instead:

pip install -r requirements.txt

Using the package

Minimal inference example (Whisper)

import numpy as np

from jaxcld import ASRModel, CVXNNLangDetectHead, NNLangDetectHead, SVMLangDetectHead

# 1) Load the base ASR model
languages = ["en", "hi", "id", "ms", "zh"]
asr = ASRModel.from_pretrained("openai/whisper-small", config={"languages": languages})

# 2) Load a language detection head artifact (choose ONE)
# head = CVXNNLangDetectHead.load("path/to/whisper-small_trained_cvx_mlp.pkl", asr)
# head = NNLangDetectHead.load("path/to/openai_whisper-small_nn_head.pkl", asr)
# head = SVMLangDetectHead.load("path/to/openai_whisper-small_linear_svm.pkl", asr)

# 3) Attach head and run inference
asr.set_lang_detect_head(head)

audio_16k_mono: np.ndarray = ...  # shape (T,), sampling rate 16kHz
pred_langs, pred_texts = asr.predict(audio_16k_mono)
print(pred_langs[0], pred_texts[0])

Pre-trained models

Trained convex heads are published on the Hugging Face Hub:

Model	Backbone	Languages	Det. Acc	WER ↓	CER ↓	HF Hub
cld-whisper-small-5lang	whisper-small	en/hi/id/ms/zh	0.98	48.23	27.47	🤗
cld-whisper-large-v3-5lang	whisper-large-v3	en/hi/id/ms/zh	0.98	31.11	19.81	🤗
cld-mms-1b-5lang	mms-1b-all	en/hi/id/ms/zh	0.96	48.10	23.47	🤗
cld-whisper-small-enzh (100–10000 samples/class)	whisper-small	en/zh	0.99–1.00	—	—	🤗

Loading from the Hub

import numpy as np
from huggingface_hub import hf_hub_download
from jaxcld import ASRModel, CVXNNLangDetectHead

languages = ["en", "hi", "id", "ms", "zh"]

# 1) Load the frozen base ASR model
asr = ASRModel.from_pretrained("openai/whisper-small", config={"languages": languages})

# 2) Download the convex head from the Hub and load it
head_path = hf_hub_download("williamhtan/cld-whisper-small-5lang", "model.pkl")
head = CVXNNLangDetectHead.load(head_path, asr)

# 3) Attach the head
asr.set_lang_detect_head(head)

# 4) Grab one clip from the dataset's test split
from datasets import load_dataset

sample = next(iter(load_dataset("williamhtan/cld-multi-dataset", split="test", streaming=True)))
audio_16k_mono = sample["audio"]["array"]   # (T,) float waveform, 16 kHz mono

# 5) Predict
pred_langs, pred_texts = asr.predict(audio_16k_mono)
print("true:", sample["lang"], "| pred:", pred_langs[0], "| text:", pred_texts[0])

For the low-resource binary model, specify the samples-per-class subfolder:

head_path = hf_hub_download("williamhtan/cld-whisper-small-enzh", "1000/model.pkl")

Citation

If you use this code in your work, please cite the paper:

@inproceedings{feng2026cld,
  title     = {Convex Low-resource Accent-Robust Language Detection in Speech Recognition},
  author    = {Feng, Miria and Tan, William and Pilanci, Mert},
  booktitle = {Proceedings of the 43rd International Conference on Machine Learning},
  year      = {2026},
  series    = {Proceedings of Machine Learning Research},
  publisher = {PMLR},
  url       = {https://icml.cc/virtual/2026/poster/64615}
}