diff --git a/transcribe.py b/transcribe.py index 1160bea..49c85db 100644 --- a/transcribe.py +++ b/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