diff --git a/servers/README.md b/servers/README.md new file mode 100644 index 0000000..5d25298 --- /dev/null +++ b/servers/README.md @@ -0,0 +1,52 @@ +# Local wrapper servers (Kyutai, Cohere) + +Minimal OpenAI-compatible HTTP servers wrapping two **local** ASR models so they +can be benchmarked through the generic `omlx` provider of `eval-transcript` +(same `POST /v1/audio/transcriptions` contract). They are referenced by the +"Kyutai STT" and "Cohere Transcribe" sections of the main `README.md`. + +Both expose: +- `GET /v1/models` → `{"data": [{"id": ...}]}` +- `POST /v1/audio/transcriptions` → `{"text": ...}` (multipart: `file`, `model`, `language`) + +| Server | Modèle | Port conseillé | Notes | +| --- | --- | ---: | --- | +| `kyutai_server.py` | `kyutai/stt-1b-en_fr` (variante `-trfs`) | 8000 | Découpage manuel ~30 s aligné sur les silences + `generate()` neuf par segment (la fenêtre de contexte de Kyutai s'effondre au-delà de ~4 min en passe continue). | +| `cohere_server.py` | `CohereLabs/cohere-transcribe-03-2026` | 8001 | Long-form géré par le modèle (le feature extractor découpe, `decode` réassemble via `audio_chunk_index`). | + +Les deux servent sur des **ports distincts** → ils peuvent tourner simultanément. + +## Lancer + +Dépendances : voir `pyproject.toml` (transformers ≥ 5.3 — requis par la classe +native Cohere ; validé sur 5.9 —, torch, librosa, soundfile, sentencepiece, +fastapi, uvicorn, python-multipart). + +```bash +# Kyutai (port 8000) +uv run uvicorn kyutai_server:app --host 127.0.0.1 --port 8000 + +# Cohere (port 8001) +uv run uvicorn cohere_server:app --host 127.0.0.1 --port 8001 +``` + +Puis transcrire via le provider `omlx` (pointer `OMLX_BASE_URL` sur le bon port) : + +```bash +OMLX_BASE_URL="http://127.0.0.1:8000/v1" uv run eval-transcript omlx transcribe \ + data/audio/.mp3 --model kyutai/stt-1b-en_fr --language fr --save + +OMLX_BASE_URL="http://127.0.0.1:8001/v1" uv run eval-transcript omlx transcribe \ + data/audio/.mp3 --model cohere-transcribe-03-2026 --language fr --save +``` + +## Note sur le chargement de Cohere + +`cohere_server.py` charge le modèle via la **classe native** +`CohereAsrForConditionalGeneration` (intégrée nativement à `transformers ≥ 5.x`), +et **non** via le chemin remote-code (`AutoModelForSpeechSeq2Seq` + +`trust_remote_code=True`). Sous transformers 5.9, le chemin remote-code applique +mal la `generation_config` (`decoder_start_token_id`) : le modèle génère alors +du texte multilingue aberrant **en ignorant l'audio**. La classe native applique +correctement la config et produit la transcription attendue. `float32` reste +obligatoire (le masque d'attention à `-1e9` déborde la plage de `float16`). diff --git a/servers/cohere_server.py b/servers/cohere_server.py new file mode 100644 index 0000000..2180929 --- /dev/null +++ b/servers/cohere_server.py @@ -0,0 +1,138 @@ +"""Serveur HTTP OpenAI-compatible minimal exposant Cohere Transcribe. + +Runtime : transformers / PyTorch (MPS sur Apple Silicon). Le modèle +`CohereLabs/cohere-transcribe-03-2026` (model_type `cohere_asr`) se charge via +la classe NATIVE `CohereAsrForConditionalGeneration` (transformers >= 5.3), et +NON via le remote-code (`AutoModelForSpeechSeq2Seq` + `trust_remote_code`), dont +le chargement est cassé sous transformers 5.9 (cf. note dans `get_model`). + +Particularité vs Kyutai : Cohere gère le long-form TOUT SEUL. Le feature +extractor découpe l'audio au-delà de `max_audio_clip_s`, et `processor.decode` +réassemble les chunks via `audio_chunk_index`. Donc PAS de découpage manuel ici. + +On le sert sur un port distinct (8001 par défaut) pour coexister avec le +serveur Kyutai (8000) ; pointer `OMLX_BASE_URL=http://127.0.0.1:8001/v1` côté +eval-transcript pour scorer Cohere. + +Endpoints attendus par le provider `omlx` de eval-transcript : + - GET /v1/models -> {"data": [{"id": ...}]} + - POST /v1/audio/transcriptions -> {"text": ...} (multipart: file, model, language) + +Lancer : uv run uvicorn cohere_server:app --host 127.0.0.1 --port 8001 +""" +from __future__ import annotations + +import os + +# Repli CPU op-par-op pour les ops non implémentées sur MPS. AVANT import torch. +os.environ.setdefault("PYTORCH_ENABLE_MPS_FALLBACK", "1") + +import tempfile +import threading +from pathlib import Path + +import librosa +import torch +from fastapi import FastAPI, File, Form, UploadFile +from transformers import AutoProcessor, CohereAsrForConditionalGeneration + +MODEL_REPO = "CohereLabs/cohere-transcribe-03-2026" +SERVED_NAME = "cohere-transcribe-03-2026" +DEVICE = "mps" if torch.backends.mps.is_available() else "cpu" +# Cohere attend de l'audio à 16 kHz (cf. model card). +SR = 16000 +# Langue par défaut des transcriptions (notre corpus est FR). +DEFAULT_LANGUAGE = "fr" +# Tokens générés par chunk interne (le modèle re-découpe le long-form lui-même). +MAX_NEW_TOKENS = 256 + +app = FastAPI(title="cohere-server") +_model = None +_processor = None +# Sérialise le lazy-load : deux requêtes concurrentes pendant le chargement +# initial chargeraient sinon le modèle 2x (pic mémoire → risque d'OOM). +_load_lock = threading.Lock() + + +def get_model(): + global _model, _processor + if _model is not None: + return _model, _processor + with _load_lock: + if _model is None: + # `cohere_asr` est intégré nativement à transformers >= 5.x : on + # utilise la classe native plutôt que le remote-code + # (`trust_remote_code` + AutoModelForSpeechSeq2Seq), dont le chemin + # de chargement est cassé sous transformers 5.9 + # (decoder_start_token_id mal appliqué → génération multilingue + # aberrante en ignorant l'audio). La classe native applique + # correctement la generation_config. + _processor = AutoProcessor.from_pretrained(MODEL_REPO) + # float32 obligatoire : le code Cohere masque l'attention avec -1e9, + # qui déborde la plage du float16 (max ~65504) → "value cannot be + # converted to type c10::Half without overflow". Plus lent mais correct. + _model = CohereAsrForConditionalGeneration.from_pretrained( + MODEL_REPO, dtype=torch.float32 + ).to(DEVICE) + _model.eval() + return _model, _processor + + +def transcribe_path(path: str, language: str) -> str: + audio, _ = librosa.load(path, sr=SR, mono=True) + model, processor = get_model() + # Le feature extractor découpe lui-même si l'audio dépasse max_audio_clip_s. + inputs = processor( + audio=audio, sampling_rate=SR, return_tensors="pt", language=language + ) + audio_chunk_index = inputs.get("audio_chunk_index") + inputs = inputs.to(model.device, dtype=model.dtype) + with torch.no_grad(): + outputs = model.generate(**inputs, max_new_tokens=MAX_NEW_TOKENS) + # Rapatrier sur CPU avant decode (certains décodeurs HF gèrent mal MPS). + outputs = outputs.to("cpu") + # `language` est obligatoire au decode dès qu'un audio_chunk_index est + # présent (réassemblage long-form) ; on le passe systématiquement. + decoded = processor.decode( + outputs, + skip_special_tokens=True, + audio_chunk_index=audio_chunk_index, + language=language, + ) + # decode renvoie une liste (un élément par item du batch). + text = decoded[0] if isinstance(decoded, (list, tuple)) else decoded + return text.strip() + + +@app.get("/v1/models") +def list_models(): + return { + "object": "list", + "data": [{"id": SERVED_NAME, "object": "model", "owned_by": "cohere"}], + } + + +# Route SYNCHRONE : `transcribe_path` est bloquante (I/O + inférence). En `def`, +# FastAPI l'exécute dans un threadpool → n'asphyxie pas l'event loop. +@app.post("/v1/audio/transcriptions") +def transcribe( + file: UploadFile = File(...), + # `model` et `response_format` sont acceptés pour la compat OpenAI mais + # ignorés : ce serveur n'expose qu'un seul modèle et ne renvoie que du texte. + model: str | None = Form(None), + language: str | None = Form(None), + response_format: str | None = Form(None), +): + suffix = Path(file.filename or "audio.wav").suffix or ".wav" + tmp = tempfile.NamedTemporaryFile(delete=False, suffix=suffix) + tmp_path = tmp.name + try: + with tmp: + tmp.write(file.file.read()) + text = transcribe_path(tmp_path, language or DEFAULT_LANGUAGE) + finally: + try: + os.unlink(tmp_path) + except OSError: + pass + return {"text": text} diff --git a/servers/kyutai_server.py b/servers/kyutai_server.py new file mode 100644 index 0000000..f2b4856 --- /dev/null +++ b/servers/kyutai_server.py @@ -0,0 +1,147 @@ +"""Serveur HTTP OpenAI-compatible minimal exposant Kyutai STT. + +Runtime : transformers / PyTorch (MPS sur Apple Silicon), via la variante +`-trfs` du modèle (classes KyutaiSpeechToText* de transformers >= 4.53). +On évite l'API streaming de moshi : la variante transformers se charge et +s'inférence comme un modèle HF classique. + +Endpoints attendus par le provider `omlx` de eval-transcript : + - GET /v1/models -> {"data": [{"id": ...}]} + - POST /v1/audio/transcriptions -> {"text": ...} (multipart: file, model, language) + +Lancer : uv run uvicorn kyutai_server:app --host 127.0.0.1 --port 8000 +""" +from __future__ import annotations + +import os + +# Certaines ops du modèle ne sont pas implémentées sur MPS : on autorise le +# repli CPU op-par-op plutôt que de planter. À définir AVANT l'import de torch. +os.environ.setdefault("PYTORCH_ENABLE_MPS_FALLBACK", "1") + +import tempfile +import threading +from pathlib import Path + +import librosa +import torch +from fastapi import FastAPI, File, Form, UploadFile +from transformers import ( + KyutaiSpeechToTextForConditionalGeneration, + KyutaiSpeechToTextProcessor, +) + +MODEL_REPO = "kyutai/stt-1b-en_fr-trfs" +SERVED_NAME = "kyutai/stt-1b-en_fr" +DEVICE = "mps" if torch.backends.mps.is_available() else "cpu" +# La variante transformers attend de l'audio à 24 kHz. +SR = 24000 +# Découpage manuel SANS overlap aligné sur les silences, comme parakeet-server : +# borne la conso mémoire/temps de `generate` sur des discours longs (jusqu'à +# ~30 min) et évite les coupures en plein mot. Kyutai étant un modèle de +# streaming, chaque segment isolé se transcrit proprement. +SEGMENT_S = 30.0 +MIN_SEGMENT_S = 0.3 + +app = FastAPI(title="kyutai-server") +_model = None +_processor = None +# Sérialise le lazy-load : deux requêtes concurrentes pendant le chargement +# initial chargeraient sinon le modèle 2x (pic mémoire → risque d'OOM). +_load_lock = threading.Lock() + + +def get_model(): + global _model, _processor + if _model is not None: + return _model, _processor + with _load_lock: + if _model is None: + _processor = KyutaiSpeechToTextProcessor.from_pretrained(MODEL_REPO) + # float16 sur MPS : ~2x plus rapide que float32, sortie identique sur + # nos tests FR. Repli float32 si MPS indispo (CPU). + dtype = torch.float16 if DEVICE == "mps" else torch.float32 + _model = KyutaiSpeechToTextForConditionalGeneration.from_pretrained( + MODEL_REPO, torch_dtype=dtype + ).to(DEVICE) + _model.eval() + return _model, _processor + + +def _silence_aligned_cuts(audio) -> list[int]: + """Bornes de coupe (~30 s) alignées sur le passage le plus calme proche.""" + seg = int(SEGMENT_S * SR) + cuts = [0] + pos = seg + while pos < len(audio): + window = audio[max(0, pos - 2 * SR):min(len(audio), pos + 2 * SR)] + if len(window): + quietest = int(librosa.feature.rms(y=window, frame_length=1024, hop_length=512).argmin()) + offset = quietest * 512 - len(window) // 2 + pos = pos + offset + cuts.append(pos) + pos += seg + cuts.append(len(audio)) + return cuts + + +def transcribe_path(path: str) -> str: + audio, _ = librosa.load(path, sr=SR, mono=True) + model, processor = get_model() + min_seg = int(MIN_SEGMENT_S * SR) + cuts = _silence_aligned_cuts(audio) + + parts: list[str] = [] + for start, end in zip(cuts, cuts[1:]): + chunk = audio[start:end] + if len(chunk) < min_seg: + continue + inputs = processor(audio=chunk, sampling_rate=SR, return_tensors="pt") + # NE PAS caster au dtype du modèle : Kyutai charge en float16 mais garde + # des biais conv en float32 ; des inputs float32 matchent ces couches. + # Forcer float16 ici casse l'encodeur ("Input type (c10::Half) and bias + # type (float) should be the same"). On laisse le dtype d'origine. + inputs = inputs.to(DEVICE) + with torch.no_grad(): + output_tokens = model.generate(**inputs) + # Rapatrier sur CPU avant decode (certains décodeurs HF gèrent mal MPS). + output_tokens = output_tokens.to("cpu") + text = processor.batch_decode(output_tokens, skip_special_tokens=True)[0].strip() + if text: + parts.append(text) + return " ".join(parts) + + +@app.get("/v1/models") +def list_models(): + return { + "object": "list", + "data": [{"id": SERVED_NAME, "object": "model", "owned_by": "kyutai"}], + } + + +# Route SYNCHRONE : `transcribe_path` est bloquante (I/O + inférence). En `def`, +# FastAPI l'exécute dans un threadpool → n'asphyxie pas l'event loop. +@app.post("/v1/audio/transcriptions") +def transcribe( + file: UploadFile = File(...), + # `model`, `language` et `response_format` sont acceptés pour la compat + # OpenAI mais ignorés : un seul modèle servi, langue gérée par le modèle, + # sortie texte uniquement. + model: str | None = Form(None), + language: str | None = Form(None), + response_format: str | None = Form(None), +): + suffix = Path(file.filename or "audio.wav").suffix or ".wav" + tmp = tempfile.NamedTemporaryFile(delete=False, suffix=suffix) + tmp_path = tmp.name + try: + with tmp: + tmp.write(file.file.read()) + text = transcribe_path(tmp_path) + finally: + try: + os.unlink(tmp_path) + except OSError: + pass + return {"text": text} diff --git a/servers/pyproject.toml b/servers/pyproject.toml new file mode 100644 index 0000000..5ccaacd --- /dev/null +++ b/servers/pyproject.toml @@ -0,0 +1,18 @@ +[project] +name = "eval-transcript-servers" +version = "0.1.0" +description = "Minimal OpenAI-compatible transcription servers wrapping local ASR models (Kyutai STT, Cohere Transcribe) for eval-transcript (transformers/PyTorch, Apple Silicon MPS)." +requires-python = ">=3.12,<3.14" +dependencies = [ + # >=5.3 requis : la classe native CohereAsrForConditionalGeneration + # (cohere_server.py) n'existe que depuis transformers 5.x ; couvre aussi + # les classes Kyutai (>=4.53). Validé sur 5.9. + "transformers>=5.3.0", + "torch", + "librosa", + "soundfile", + "sentencepiece", + "fastapi", + "uvicorn[standard]", + "python-multipart", +]