112 lines · 3.4 KB
python
| 1 | """Hugging Face Whisper — the ear. |
| 2 | |
| 3 | English Spaces: openai/whisper-small.en or medium.en. |
| 4 | Mixed language: openai/whisper-large-v3. |
| 5 | |
| 6 | Long Spaces should be chunked with ffmpeg (see render.wav_for_asr) so a CPU |
| 7 | box doesn't melt. Add the chunk offset back onto every timestamp. |
| 8 | """ |
| 9 | |
| 10 | from __future__ import annotations |
| 11 | |
| 12 | from pathlib import Path |
| 13 | |
| 14 | DEFAULT_MODEL = "openai/whisper-small.en" |
| 15 | |
| 16 | |
| 17 | def transcribe( |
| 18 | wav: str | Path, |
| 19 | *, |
| 20 | model: str = DEFAULT_MODEL, |
| 21 | word_timestamps: bool = False, |
| 22 | device: str | None = None, |
| 23 | chunk_offset: float = 0.0, |
| 24 | ) -> list[dict]: |
| 25 | """Return [{start, end, text}] (and optional `words`).""" |
| 26 | asr = _pipeline(model=model, word_timestamps=word_timestamps, device=device) |
| 27 | kwargs: dict = {} |
| 28 | if word_timestamps: |
| 29 | kwargs["return_timestamps"] = "word" |
| 30 | out = asr(str(wav), **kwargs) |
| 31 | return normalize_asr(out, chunk_offset=chunk_offset) |
| 32 | |
| 33 | |
| 34 | def transcribe_chunks( |
| 35 | wavs: list[tuple[Path, float]], |
| 36 | **kwargs, |
| 37 | ) -> list[dict]: |
| 38 | """wavs is [(path, offset_seconds), ...].""" |
| 39 | cues: list[dict] = [] |
| 40 | for path, offset in wavs: |
| 41 | cues.extend(transcribe(path, chunk_offset=offset, **kwargs)) |
| 42 | return cues |
| 43 | |
| 44 | |
| 45 | def normalize_asr(out: object, chunk_offset: float = 0.0) -> list[dict]: |
| 46 | if not isinstance(out, dict): |
| 47 | text = str(out or "").strip() |
| 48 | return [{"start": chunk_offset, "end": chunk_offset, "text": text}] if text else [] |
| 49 | chunks = out.get("chunks") or [] |
| 50 | cues: list[dict] = [] |
| 51 | for chunk in chunks: |
| 52 | if not isinstance(chunk, dict): |
| 53 | continue |
| 54 | ts = chunk.get("timestamp") or (None, None) |
| 55 | if not isinstance(ts, (tuple, list)): |
| 56 | ts = (None, None) |
| 57 | start_raw, end_raw = (ts + (None, None))[:2] |
| 58 | text = (chunk.get("text") or "").strip() |
| 59 | if not text: |
| 60 | continue |
| 61 | start = float(start_raw or 0) + chunk_offset |
| 62 | end = float(end_raw if end_raw is not None else start_raw or 0) + chunk_offset |
| 63 | cue: dict = {"start": start, "end": end, "text": text} |
| 64 | words = chunk.get("words") |
| 65 | if isinstance(words, list): |
| 66 | cue["words"] = words |
| 67 | cues.append(cue) |
| 68 | if cues: |
| 69 | return cues |
| 70 | text = (out.get("text") or "").strip() |
| 71 | if not text: |
| 72 | return [] |
| 73 | return [{"start": chunk_offset, "end": chunk_offset, "text": text}] |
| 74 | |
| 75 | |
| 76 | def _pipeline(*, model: str, word_timestamps: bool, device: str | None): |
| 77 | try: |
| 78 | from transformers import pipeline |
| 79 | except ImportError as exc: |
| 80 | raise RuntimeError( |
| 81 | "transformers is required for Whisper. " |
| 82 | "pip install space-tape (or: pip install transformers torch)" |
| 83 | ) from exc |
| 84 | try: |
| 85 | import torch # noqa: F401 |
| 86 | except ImportError as exc: |
| 87 | raise RuntimeError( |
| 88 | "torch is required for Whisper. A CPU wheel is fine:\n" |
| 89 | " pip install torch --index-url https://download.pytorch.org/whl/cpu" |
| 90 | ) from exc |
| 91 | |
| 92 | resolved_device = device if device is not None else _default_device() |
| 93 | ts: bool | str = "word" if word_timestamps else True |
| 94 | return pipeline( |
| 95 | "automatic-speech-recognition", |
| 96 | model=model, |
| 97 | chunk_length_s=30, |
| 98 | return_timestamps=ts, |
| 99 | device=resolved_device, |
| 100 | ) |
| 101 | |
| 102 | |
| 103 | def _default_device() -> str | int: |
| 104 | try: |
| 105 | import torch |
| 106 | |
| 107 | if torch.cuda.is_available(): |
| 108 | return 0 |
| 109 | except Exception: |
| 110 | pass |
| 111 | return "cpu" |