Add local translation using NLLB-200 (MPS accelerated)

transcribe.py

@@ -5,12 +5,15 @@ import queue
 import sys
 import torch
 from silero_vad import load_silero_vad, get_speech_timestamps
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
 
 # Parameters
-MODEL_PATH = "mlx-community/whisper-small.en-mlx" # MLX optimized small model
+WHISPER_MODEL = "mlx-community/whisper-small.en-mlx"
+TRANSLATE_MODEL = "facebook/nllb-200-distilled-600M"
+TARGET_LANG = "spa_Latn" # Spanish (Latin America) - change as needed
 CHANNELS = 1
 SAMPLERATE = 16000
-BLOCK_SIZE = 512 # Silero VAD prefers 512, 1024, or 1536
+BLOCK_SIZE = 512
 VAD_THRESHOLD = 0.5 
 BUFFER_LIMIT = SAMPLERATE * 30 
 MIN_SILENCE_DURATION_MS = 500 
@@ -23,15 +26,25 @@ def callback(indata, frames, time, status):
     audio_queue.put(indata.copy())
 
 def main():
-    print(f"Loading MLX-optimized Whisper model '{MODEL_PATH}'...")
-    # mlx-whisper uses the same model names or Hugging Face paths
+    # Set device for translation model (using MPS for Mac M chips)
+    device = "mps" if torch.backends.mps.is_available() else "cpu"
+    print(f"Using device: {device}")
+
+    print(f"Loading Whisper model '{WHISPER_MODEL}'...")
+    # Whisper MLX runs on its own optimized path
     
+    print(f"Loading translation model '{TRANSLATE_MODEL}'...")
+    tokenizer = AutoTokenizer.from_pretrained(TRANSLATE_MODEL)
+    model = AutoModelForSeq2SeqLM.from_pretrained(TRANSLATE_MODEL).to(device)
+    translator = pipeline("translation", model=model, tokenizer=tokenizer, 
+                          src_lang="eng_Latn", tgt_lang=TARGET_LANG, device=device)
+
     print("Loading Silero VAD model...")
     vad_model = load_silero_vad()
     
     print("Models loaded.")
 
-    print("\nStarting live transcription (MLX + VAD)... (Press Ctrl+C to stop)")
+    print(f"\nStarting live transcription & translation to {TARGET_LANG}... (Press Ctrl+C to stop)")
     
     audio_buffer = []
     speech_started = False
@@ -62,12 +75,17 @@ def main():
                         
                         if (buffer_len_samples - last_end) > (SAMPLERATE * MIN_SILENCE_DURATION_MS / 1000) or buffer_len_samples > BUFFER_LIMIT:
                             
-                            # Transcribe with MLX
-                            result = mlx_whisper.transcribe(current_audio, path_or_hf_repo=MODEL_PATH)
-                            text = result['text'].strip()
+                            # 1. Transcribe (English)
+                            result = mlx_whisper.transcribe(current_audio, path_or_hf_repo=WHISPER_MODEL)
+                            original_text = result['text'].strip()
                             
-                            if text:
-                                print(f"Transcription: {text}")
+                            if original_text:
+                                # 2. Translate
+                                translation = translator(original_text)
+                                translated_text = translation[0]['translation_text']
+                                
+                                print(f"\nEN: {original_text}")
+                                print(f"ES: {translated_text}")
                             
                             audio_buffer = []
                             speech_started = False