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

Support multiple simultaneous translation languages (ES, FR, AR)

Browse Source
This commit is contained in:
Adolfo Reyna
2026-02-26 21:18:31 -05:00
parent c4a39ef29e
commit 3166d6e7d5
1 changed files with 26 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files

42
transcribe.py
View File

@@ -6,7 +6,6 @@ try:
import lzma
except ImportError:
mock_lzma = MagicMock()
# Add common constants that libraries expect from lzma
mock_lzma.FORMAT_XZ = 1
mock_lzma.FORMAT_ALONE = 2
mock_lzma.FORMAT_RAW = 3
@@ -27,7 +26,13 @@ from transformers import MarianMTModel, MarianTokenizer
# Parameters
WHISPER_MODEL = "mlx-community/whisper-small.en-mlx"
TRANSLATE_MODEL = "Helsinki-NLP/opus-mt-en-es"
# List of language pairs (English to ...)
TARGET_LANGS = {
"Spanish": "Helsinki-NLP/opus-mt-en-es",
"French": "Helsinki-NLP/opus-mt-en-fr",
"Arabic": "Helsinki-NLP/opus-mt-en-ar"
}
CHANNELS = 1
SAMPLERATE = 16000
BLOCK_SIZE = 512
@@ -43,22 +48,26 @@ def callback(indata, frames, time, status):
audio_queue.put(indata.copy())
def main():
# Set device for translation model
device = "mps" if torch.backends.mps.is_available() else "cpu"
print(f"Using device for translation: {device}")
print(f"Using device: {device}")
print(f"Loading Whisper model '{WHISPER_MODEL}'...")
print(f"Loading dedicated translation model '{TRANSLATE_MODEL}'...")
tokenizer = MarianTokenizer.from_pretrained(TRANSLATE_MODEL)
model = MarianMTModel.from_pretrained(TRANSLATE_MODEL).to(device)
# Dictionary to hold models and tokenizers
translation_engines = {}
for lang_name, model_id in TARGET_LANGS.items():
print(f"Loading {lang_name} translation model ({model_id})...")
tokenizer = MarianTokenizer.from_pretrained(model_id)
model = MarianMTModel.from_pretrained(model_id).to(device)
translation_engines[lang_name] = (model, tokenizer)
print("Loading Silero VAD model...")
vad_model = load_silero_vad()
print("Models loaded.")
print(f"\nStarting live transcription & translation... (Press Ctrl+C to stop)")
print(f"\nStarting live transcription & multiple translations... (Press Ctrl+C to stop)")
audio_buffer = []
speech_started = False
@@ -89,19 +98,20 @@ def main():
if (buffer_len_samples - last_end) > (SAMPLERATE * MIN_SILENCE_DURATION_MS / 1000) or buffer_len_samples > BUFFER_LIMIT:
# 1. Transcribe (English)
# 1. Transcribe once
result = mlx_whisper.transcribe(current_audio, path_or_hf_repo=WHISPER_MODEL)
original_text = result['text'].strip()
if original_text:
# 2. Translate
inputs = tokenizer(original_text, return_tensors="pt", padding=True).to(device)
with torch.no_grad():
translated_tokens = model.generate(**inputs)
translated_text = tokenizer.decode(translated_tokens[0], skip_special_tokens=True)
print(f"\n[EN]: {original_text}")
print(f"\nEN: {original_text}")
print(f"ES: {translated_text}")
# 2. Translate to all targets
for lang_name, (model, tokenizer) in translation_engines.items():
inputs = tokenizer(original_text, return_tensors="pt", padding=True).to(device)
with torch.no_grad():
translated_tokens = model.generate(**inputs)
translated_text = tokenizer.decode(translated_tokens[0], skip_special_tokens=True)
print(f"[{lang_name[:2].upper()}]: {translated_text}")
audio_buffer = []
speech_started = False