Add mlx transcriber using mlx-whisper and pyannote
Self-contained uv script handles its own dependencies. Runs
pyannote on MPS when available to avoid CPU bottleneck.

Assisted-by: Claude Opus 4.6 via pi
change ynnqzllrmkuwqwurxwznwpvvylosqkmk
commit ef2189f374657672fa7cf034a70797700bfb0d90
author Alpha Chen <alpha@kejadlen.dev>
date
parent qrqwwuwv
diff --git a/Rakefile b/Rakefile
index 1c16d92..c57222b 100644
--- a/Rakefile
+++ b/Rakefile
@@ -25,7 +25,7 @@ end
 
 load File.expand_path("transcribers.rake", __dir__)
 
-TRANSCRIBER = Transcribers.resolve(ENV.fetch("TRANSCRIBER", "whisperx"))
+TRANSCRIBER = Transcribers.resolve(ENV.fetch("TRANSCRIBER", "whisper-cpp-large"))
 TRANSCRIBER.register
 
 Rake::Task[HRN_FEED.to_s].invoke
diff --git a/bin/mlx-transcribe b/bin/mlx-transcribe
new file mode 100755
index 0000000..fd46548
--- /dev/null
+++ b/bin/mlx-transcribe
@@ -0,0 +1,131 @@
+#!/usr/bin/env -S uv run --script
+# /// script
+# requires-python = ">=3.10"
+# dependencies = [
+#     "mlx-whisper",
+#     "pyannote.audio",
+#     "torch",
+#     "torchaudio",
+# ]
+# ///
+
+"""Transcribe audio with mlx-whisper and diarize with pyannote."""
+
+import argparse
+import sys
+from pathlib import Path
+
+import mlx_whisper
+from pyannote.audio import Pipeline
+
+
+def transcribe(audio_path: str, model: str) -> dict:
+    """Run mlx-whisper transcription, returning segments with timestamps."""
+    return mlx_whisper.transcribe(
+        audio_path,
+        path_or_hf_repo=model,
+        word_timestamps=True,
+        verbose=False,
+    )
+
+
+def diarize(audio_path: str, hf_token: str) -> Pipeline:
+    """Run pyannote speaker diarization."""
+    import torch
+
+    pipeline = Pipeline.from_pretrained(
+        "pyannote/speaker-diarization-3.1",
+        token=hf_token,
+    )
+    if torch.backends.mps.is_available():
+        pipeline.to(torch.device("mps"))
+    elif torch.cuda.is_available():
+        pipeline.to(torch.device("cuda"))
+    return pipeline(audio_path)
+
+
+def assign_speakers(segments: list[dict], diarization) -> list[dict]:
+    """Assign a speaker label to each whisper segment based on diarization overlap."""
+    labeled = []
+    for seg in segments:
+        seg_start = seg["start"]
+        seg_end = seg["end"]
+        # Find the diarization speaker with the most overlap.
+        best_speaker = None
+        best_overlap = 0.0
+        for turn, _, speaker in diarization.itertracks(yield_label=True):
+            overlap_start = max(seg_start, turn.start)
+            overlap_end = min(seg_end, turn.end)
+            overlap = max(0.0, overlap_end - overlap_start)
+            if overlap > best_overlap:
+                best_overlap = overlap
+                best_speaker = speaker
+        labeled.append({**seg, "speaker": best_speaker or "UNKNOWN"})
+    return labeled
+
+
+def format_time(seconds: float) -> str:
+    """Format seconds as H:MM:SS or M:SS."""
+    h = int(seconds) // 3600
+    m = (int(seconds) % 3600) // 60
+    s = int(seconds) % 60
+    if h > 0:
+        return f"{h}:{m:02d}:{s:02d}"
+    return f"{m}:{s:02d}"
+
+
+def format_output(segments: list[dict]) -> str:
+    """Format labeled segments as timestamped, speaker-attributed lines."""
+    lines = []
+    for seg in segments:
+        start = format_time(seg["start"])
+        end = format_time(seg["end"])
+        speaker = seg["speaker"]
+        text = seg["text"].strip()
+        lines.append(f"[{start} → {end} | {speaker}] {text}")
+    return "\n".join(lines)
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("audio", help="Path to audio file")
+    parser.add_argument("output", help="Path to write transcript")
+    parser.add_argument(
+        "--model",
+        default="mlx-community/whisper-large-v3-turbo",
+        help="MLX whisper model (default: mlx-community/whisper-large-v3-turbo)",
+    )
+    parser.add_argument(
+        "--hf-token",
+        default=None,
+        help="Hugging Face token for pyannote (default: $HUGGING_FACE_TOKEN)",
+    )
+    args = parser.parse_args()
+
+    import os
+
+    hf_token = args.hf_token or os.environ.get("HUGGING_FACE_TOKEN")
+    if not hf_token:
+        print("Set --hf-token or $HUGGING_FACE_TOKEN for diarization.", file=sys.stderr)
+        sys.exit(1)
+
+    audio = args.audio
+    if not Path(audio).exists():
+        print(f"Audio file not found: {audio}", file=sys.stderr)
+        sys.exit(1)
+
+    print(f"Transcribing {audio}...", file=sys.stderr)
+    result = transcribe(audio, args.model)
+
+    print("Diarizing...", file=sys.stderr)
+    diarization = diarize(audio, hf_token)
+
+    segments = assign_speakers(result["segments"], diarization)
+    text = format_output(segments)
+
+    Path(args.output).write_text(text, encoding="utf-8")
+    print(f"Done. Wrote {len(segments)} segments to {args.output}.", file=sys.stderr)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/transcribers.rake b/transcribers.rake
index aca6a8a..9727922 100644
--- a/transcribers.rake
+++ b/transcribers.rake
@@ -9,8 +9,9 @@ module Transcribers
     when "whisper-cpp-large" then WhisperCppLarge.new
     when "whisper-cpp-tdrz" then WhisperCppTdrz.new
     when "sous_chef" then SousChef.new
+    when "mlx" then Mlx.new
     else
-      abort "Unknown transcriber: #{name}. Use 'whisperx', 'whisper-cpp-large', 'whisper-cpp-tdrz', or 'sous_chef'."
+      abort "Unknown transcriber: #{name}. Use 'whisperx', 'whisper-cpp-large', 'whisper-cpp-tdrz', 'sous_chef', or 'mlx'."
     end
   end
 
@@ -64,8 +65,7 @@ module Transcribers
   end
 
   class WhisperCppLarge < WhisperCpp
-    # MODEL = "large-v3-turbo"
-    MODEL = "large-v3"
+    MODEL = "large-v3-turbo"
 
     def name = "whisper-cpp-large"
     def prereqs = [Transcribers.model_path(MODEL).to_s]
@@ -120,4 +120,16 @@ module Transcribers
       sh BINARY.to_s, audio_path, transcript_path
     end
   end
+
+  class Mlx < Base
+    SCRIPT = Pathname("bin/mlx-transcribe")
+
+    def name = "mlx"
+    def register; end
+
+    def call(audio_path, transcript_path)
+      hf_token = ENV.fetch("HUGGING_FACE_TOKEN") { abort "Set HUGGING_FACE_TOKEN for diarization." }
+      sh SCRIPT.to_s, audio_path, transcript_path, "--hf-token", hf_token
+    end
+  end
 end