diff --git a/transcribe.py b/transcribe.py index d501687..f8c7d42 100644 --- a/transcribe.py +++ b/transcribe.py @@ -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