adolforeyna/whisper-translation
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add local translation using NLLB-200 (MPS accelerated)

Browse Source
This commit is contained in:
Adolfo Reyna
2026-02-26 21:08:08 -05:00
parent 3ce95633fc
commit 035d6e9358
1 changed files with 28 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

38
transcribe.py
View File

@@ -5,12 +5,15 @@ import queue
import sys import sys
import torch import torch
from silero_vad import load_silero_vad, get_speech_timestamps from silero_vad import load_silero_vad, get_speech_timestamps
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
# Parameters # 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 CHANNELS = 1
SAMPLERATE = 16000 SAMPLERATE = 16000
BLOCK_SIZE = 512 # Silero VAD prefers 512, 1024, or 1536 BLOCK_SIZE = 512
VAD_THRESHOLD = 0.5 VAD_THRESHOLD = 0.5
BUFFER_LIMIT = SAMPLERATE * 30 BUFFER_LIMIT = SAMPLERATE * 30
MIN_SILENCE_DURATION_MS = 500 MIN_SILENCE_DURATION_MS = 500
@@ -23,15 +26,25 @@ def callback(indata, frames, time, status):
audio_queue.put(indata.copy()) audio_queue.put(indata.copy())
def main(): def main():
print(f"Loading MLX-optimized Whisper model '{MODEL_PATH}'...") # Set device for translation model (using MPS for Mac M chips)
# mlx-whisper uses the same model names or Hugging Face paths 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...") print("Loading Silero VAD model...")
vad_model = load_silero_vad() vad_model = load_silero_vad()
print("Models loaded.") 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 = [] audio_buffer = []
speech_started = False 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: if (buffer_len_samples - last_end) > (SAMPLERATE * MIN_SILENCE_DURATION_MS / 1000) or buffer_len_samples > BUFFER_LIMIT:
# Transcribe with MLX # 1. Transcribe (English)
result = mlx_whisper.transcribe(current_audio, path_or_hf_repo=MODEL_PATH) result = mlx_whisper.transcribe(current_audio, path_or_hf_repo=WHISPER_MODEL)
text = result['text'].strip() original_text = result['text'].strip()
if text: if original_text:
print(f"Transcription: {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 = [] audio_buffer = []
speech_started = False speech_started = False