whisper-slim-mlx

Versión podada (slim) de mlx-community/whisper-large-v3-mlx para Apple Silicon (MLX). Se eliminaron 25 bloques Transformer (15 encoder + 10 decoder) partiendo de un estudio de ablación bloque a bloque sobre un corpus médico en español.

	Original	Slim
Encoder blocks	32	17
Decoder blocks	32	22
Parámetros	~1.55 B	~0.97 B
weights.npz	~3.1 GB	1.97 GB
n_mels	128	128
n_vocab	51866	51866

Bloques eliminados (ver slim_metadata.json):

• Encoder: [1, 2, 3, 4, 8, 9, 10, 11, 12, 13, 14, 15, 18, 19, 21]
• Decoder: [3, 7, 8, 11, 13, 18, 24, 26, 29, 30]

⚠️ Este checkpoint es una poda one-shot agresiva SIN reentrenamiento posterior. Sirve como punto de partida para knowledge distillation o fine-tuning con LoRA. Sin recuperación, el WER colapsa — no usar en producción tal cual. Para un pipeline funcional, reentrena con destilación (teacher = whisper-large-v3) o LoRA sobre tu dominio.

---

Requisitos

• macOS con Apple Silicon (M1/M2/M3/M4)
• Python 3.10+

pip install mlx>=0.22.0 mlx-whisper>=0.4.0 huggingface_hub numpy soundfile

Descarga

from huggingface_hub import snapshot_download

model_path = snapshot_download(repo_id="MrZeggers/whisper-slim-mlx")
print(model_path)
# ~/.cache/huggingface/hub/models--MrZeggers--whisper-slim-mlx/snapshots/<hash>

O con la CLI:

hf download MrZeggers/whisper-slim-mlx --local-dir ./whisper-slim-mlx

Uso con mlx-whisper

mlx_whisper.load_models.load_model lee config.json + weights.npz directamente:

from huggingface_hub import snapshot_download
from mlx_whisper.load_models import load_model
from mlx_whisper.tokenizer import get_tokenizer
from mlx_whisper.audio import load_audio, pad_or_trim, log_mel_spectrogram
import mlx.core as mx

model_path = snapshot_download(repo_id="MrZeggers/whisper-slim-mlx")
model = load_model(model_path)

tokenizer = get_tokenizer(multilingual=True, language="es", task="transcribe")

audio = load_audio("mi_audio.wav")
audio = pad_or_trim(audio)
mel = log_mel_spectrogram(audio, n_mels=128)           # ¡128 mels!
mel = mx.expand_dims(mel, 0)                           # [1, n_frames, 128]

# Encoder
audio_features = model.encoder(mel)

# Decoder greedy
sot = [tokenizer.sot, tokenizer.language_token, tokenizer.transcribe, tokenizer.no_timestamps]
tokens = mx.array([sot])
for _ in range(224):
    logits = model.decoder(tokens, audio_features)
    next_tok = mx.argmax(logits[:, -1, :], axis=-1, keepdims=True)
    tokens = mx.concatenate([tokens, next_tok], axis=-1)
    if next_tok.item() == tokenizer.eot:
        break

text = tokenizer.decode(tokens[0].tolist()[len(sot):-1])
print(text)

Estructura del repositorio

whisper-slim-mlx/
├── config.json            # dims MLX (n_audio_layer=17, n_text_layer=22, n_mels=128)
├── weights.npz            # pesos remapeados a los nuevos índices de bloque
├── slim_metadata.json     # qué bloques se mantuvieron/eliminaron
└── README.md

El remapeo de claves en weights.npz re-numera los bloques retenidos a índices contiguos (0..n-1), de modo que load_model puede cargarlos sin modificaciones. alignment_heads también se recalcula para el nuevo número de capas de decoder.

Aplicar LoRA sobre este modelo

from mlx_whisper.load_models import load_model
from train_mlx import apply_lora, load_checkpoint  # tu script

model = load_model(model_path)
model = apply_lora(model)                 # Q/V en cada bloque superviviente
load_checkpoint(model, "ruta/al/checkpoint-LoRA")

El apply_lora debe recorrer los bloques por índice del modelo slim (0..16 en encoder, 0..21 en decoder), no por el índice original.

Origen metodológico

Poda derivada del estudio de ablación individual bloque a bloque (ablacion_capas.py) sobre un set de 66 audios médicos en español. Se eliminaron los bloques con delta_WER < 0.5pp al ser anulados individualmente. La literatura (Distil-Whisper, Whisper Turbo, NASH) sugiere que una estrategia más robusta pasa por proteger el encoder y podar el decoder por selección uniforme, reentrenando entre rondas — este checkpoint conserva la poda cruda como baseline experimental.

Licencia

MIT — hereda la licencia del modelo base Whisper.

Citación

Si usas este modelo en un trabajo académico, cita también Whisper original:

@article{radford2022whisper,
  title={Robust Speech Recognition via Large-Scale Weak Supervision},
  author={Radford, Alec and Kim, Jong Wook and Xu, Tao and Brockman, Greg and McLeavey, Christine and Sutskever, Ilya},
  journal={arXiv:2212.04356},
  year={2022}
}