diff --git a/transcribe.py b/transcribe.py index f921f8b..c497c27 100644 --- a/transcribe.py +++ b/transcribe.py @@ -25,8 +25,9 @@ from silero_vad import load_silero_vad, get_speech_timestamps from transformers import MarianMTModel, MarianTokenizer # Parameters -WHISPER_MODEL = "mlx-community/whisper-small.en-mlx" -# List of language pairs (English to ...) +# Using the multilingual small model +WHISPER_MODEL = "mlx-community/whisper-small-mlx" +# Translation models (English -> Target) TARGET_LANGS = { "Spanish": "Helsinki-NLP/opus-mt-en-es", "French": "Helsinki-NLP/opus-mt-en-fr", @@ -51,8 +52,8 @@ def main(): device = "mps" if torch.backends.mps.is_available() else "cpu" print(f"Using device: {device}") - # 1. Load models first - print(f"Loading Whisper model '{WHISPER_MODEL}'...") + # 1. Load models + print(f"Loading Multilingual Whisper model '{WHISPER_MODEL}'...") translation_engines = {} for lang_name, model_id in TARGET_LANGS.items(): @@ -76,13 +77,12 @@ def main(): print("Invalid input, using default device.") device_index = None - print(f"\nStarting live transcription & multiple translations using device {device_index if device_index is not None else 'default'}... (Press Ctrl+C to stop)") + print(f"\nStarting Multilingual live transcription... (Press Ctrl+C to stop)") audio_buffer = [] speech_started = False try: - # 3. Start stream with sd.InputStream(samplerate=SAMPLERATE, channels=CHANNELS, callback=callback, blocksize=BLOCK_SIZE, device=device_index): while True: while not audio_queue.empty(): @@ -108,23 +108,35 @@ def main(): if (buffer_len_samples - last_end) > (SAMPLERATE * MIN_SILENCE_DURATION_MS / 1000) or buffer_len_samples > BUFFER_LIMIT: - # Transcribe once - result = mlx_whisper.transcribe(current_audio, path_or_hf_repo=WHISPER_MODEL) - original_text = result['text'].strip() + # 1. Transcribe & Detect Language + # We use task="transcribe" to get the original text and detect language + transcription_result = mlx_whisper.transcribe(current_audio, path_or_hf_repo=WHISPER_MODEL) + original_text = transcription_result['text'].strip() + detected_lang = transcription_result.get('language', 'unknown') if original_text: - if len(original_text) > 250: - original_text = original_text[:247] + "..." + print(f"\n[{detected_lang.upper()} detected]: {original_text}") - print(f"\n[EN]: {original_text}") - - # 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, max_new_tokens=150) - translated_text = tokenizer.decode(translated_tokens[0], skip_special_tokens=True) - print(f"[{lang_name[:2].upper()}]: {translated_text}") + # 2. Bridge to English if not already English + if detected_lang != "en": + # Use Whisper to translate the segment to English + bridge_result = mlx_whisper.transcribe(current_audio, path_or_hf_repo=WHISPER_MODEL, task="translate") + english_text = bridge_result['text'].strip() + print(f"[EN Bridge]: {english_text}") + else: + english_text = original_text + + # 3. Translate from English to other languages + if english_text: + if len(english_text) > 250: + english_text = english_text[:247] + "..." + + for lang_name, (model, tokenizer) in translation_engines.items(): + inputs = tokenizer(english_text, return_tensors="pt", padding=True).to(device) + with torch.no_grad(): + translated_tokens = model.generate(**inputs, max_new_tokens=150) + translated_text = tokenizer.decode(translated_tokens[0], skip_special_tokens=True) + print(f"[{lang_name[:2].upper()}]: {translated_text}") audio_buffer = [] speech_started = False