Support multilingual detection and bridge translation via Whisper small
This commit is contained in:
@@ -25,8 +25,9 @@ from silero_vad import load_silero_vad, get_speech_timestamps
|
|||||||
from transformers import MarianMTModel, MarianTokenizer
|
from transformers import MarianMTModel, MarianTokenizer
|
||||||
|
|
||||||
# Parameters
|
# Parameters
|
||||||
WHISPER_MODEL = "mlx-community/whisper-small.en-mlx"
|
# Using the multilingual small model
|
||||||
# List of language pairs (English to ...)
|
WHISPER_MODEL = "mlx-community/whisper-small-mlx"
|
||||||
|
# Translation models (English -> Target)
|
||||||
TARGET_LANGS = {
|
TARGET_LANGS = {
|
||||||
"Spanish": "Helsinki-NLP/opus-mt-en-es",
|
"Spanish": "Helsinki-NLP/opus-mt-en-es",
|
||||||
"French": "Helsinki-NLP/opus-mt-en-fr",
|
"French": "Helsinki-NLP/opus-mt-en-fr",
|
||||||
@@ -51,8 +52,8 @@ def main():
|
|||||||
device = "mps" if torch.backends.mps.is_available() else "cpu"
|
device = "mps" if torch.backends.mps.is_available() else "cpu"
|
||||||
print(f"Using device: {device}")
|
print(f"Using device: {device}")
|
||||||
|
|
||||||
# 1. Load models first
|
# 1. Load models
|
||||||
print(f"Loading Whisper model '{WHISPER_MODEL}'...")
|
print(f"Loading Multilingual Whisper model '{WHISPER_MODEL}'...")
|
||||||
|
|
||||||
translation_engines = {}
|
translation_engines = {}
|
||||||
for lang_name, model_id in TARGET_LANGS.items():
|
for lang_name, model_id in TARGET_LANGS.items():
|
||||||
@@ -76,13 +77,12 @@ def main():
|
|||||||
print("Invalid input, using default device.")
|
print("Invalid input, using default device.")
|
||||||
device_index = None
|
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 = []
|
audio_buffer = []
|
||||||
speech_started = False
|
speech_started = False
|
||||||
|
|
||||||
try:
|
try:
|
||||||
# 3. Start stream
|
|
||||||
with sd.InputStream(samplerate=SAMPLERATE, channels=CHANNELS, callback=callback, blocksize=BLOCK_SIZE, device=device_index):
|
with sd.InputStream(samplerate=SAMPLERATE, channels=CHANNELS, callback=callback, blocksize=BLOCK_SIZE, device=device_index):
|
||||||
while True:
|
while True:
|
||||||
while not audio_queue.empty():
|
while not audio_queue.empty():
|
||||||
@@ -108,19 +108,31 @@ 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 once
|
# 1. Transcribe & Detect Language
|
||||||
result = mlx_whisper.transcribe(current_audio, path_or_hf_repo=WHISPER_MODEL)
|
# We use task="transcribe" to get the original text and detect language
|
||||||
original_text = result['text'].strip()
|
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 original_text:
|
||||||
if len(original_text) > 250:
|
print(f"\n[{detected_lang.upper()} detected]: {original_text}")
|
||||||
original_text = original_text[:247] + "..."
|
|
||||||
|
|
||||||
print(f"\n[EN]: {original_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] + "..."
|
||||||
|
|
||||||
# Translate to all targets
|
|
||||||
for lang_name, (model, tokenizer) in translation_engines.items():
|
for lang_name, (model, tokenizer) in translation_engines.items():
|
||||||
inputs = tokenizer(original_text, return_tensors="pt", padding=True).to(device)
|
inputs = tokenizer(english_text, return_tensors="pt", padding=True).to(device)
|
||||||
with torch.no_grad():
|
with torch.no_grad():
|
||||||
translated_tokens = model.generate(**inputs, max_new_tokens=150)
|
translated_tokens = model.generate(**inputs, max_new_tokens=150)
|
||||||
translated_text = tokenizer.decode(translated_tokens[0], skip_special_tokens=True)
|
translated_text = tokenizer.decode(translated_tokens[0], skip_special_tokens=True)
|
||||||
|
|||||||
Reference in New Issue
Block a user