156 lines
6.3 KiB
Python
156 lines
6.3 KiB
Python
import sys
|
|
from unittest.mock import MagicMock
|
|
|
|
# Comprehensive workaround for missing _lzma in some Python builds
|
|
try:
|
|
import lzma
|
|
except ImportError:
|
|
mock_lzma = MagicMock()
|
|
mock_lzma.FORMAT_XZ = 1
|
|
mock_lzma.FORMAT_ALONE = 2
|
|
mock_lzma.FORMAT_RAW = 3
|
|
mock_lzma.CHECK_NONE = 0
|
|
mock_lzma.CHECK_CRC32 = 1
|
|
mock_lzma.CHECK_CRC64 = 4
|
|
mock_lzma.CHECK_SHA256 = 10
|
|
sys.modules["_lzma"] = MagicMock()
|
|
sys.modules["lzma"] = mock_lzma
|
|
|
|
import mlx_whisper
|
|
import numpy as np
|
|
import sounddevice as sd
|
|
import queue
|
|
import torch
|
|
from silero_vad import load_silero_vad, get_speech_timestamps
|
|
from transformers import MarianMTModel, MarianTokenizer
|
|
|
|
# Parameters
|
|
# 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",
|
|
"Arabic": "Helsinki-NLP/opus-mt-en-ar"
|
|
}
|
|
|
|
CHANNELS = 1
|
|
SAMPLERATE = 16000
|
|
BLOCK_SIZE = 512
|
|
VAD_THRESHOLD = 0.5
|
|
BUFFER_LIMIT = SAMPLERATE * 30
|
|
MIN_SILENCE_DURATION_MS = 500
|
|
|
|
audio_queue = queue.Queue()
|
|
|
|
def callback(indata, frames, time, status):
|
|
if status:
|
|
print(status, file=sys.stderr)
|
|
audio_queue.put(indata.copy())
|
|
|
|
def main():
|
|
device = "mps" if torch.backends.mps.is_available() else "cpu"
|
|
print(f"Using device: {device}")
|
|
|
|
# 1. Load models
|
|
print(f"Loading Multilingual Whisper model '{WHISPER_MODEL}'...")
|
|
|
|
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.")
|
|
|
|
# 2. Select Audio Device
|
|
print("\nAvailable Audio Devices:")
|
|
print(sd.query_devices())
|
|
|
|
try:
|
|
device_input = input("\nSelect input device index (or press Enter for default): ")
|
|
device_index = int(device_input) if device_input.strip() else None
|
|
except ValueError:
|
|
print("Invalid input, using default device.")
|
|
device_index = None
|
|
|
|
print(f"\nStarting Multilingual live transcription... (Press Ctrl+C to stop)")
|
|
|
|
audio_buffer = []
|
|
speech_started = False
|
|
|
|
try:
|
|
with sd.InputStream(samplerate=SAMPLERATE, channels=CHANNELS, callback=callback, blocksize=BLOCK_SIZE, device=device_index):
|
|
while True:
|
|
while not audio_queue.empty():
|
|
data = audio_queue.get()
|
|
audio_buffer.append(data.flatten())
|
|
|
|
if len(audio_buffer) > 0:
|
|
current_audio = np.concatenate(audio_buffer)
|
|
audio_tensor = torch.from_numpy(current_audio)
|
|
|
|
speech_timestamps = get_speech_timestamps(
|
|
audio_tensor,
|
|
vad_model,
|
|
sampling_rate=SAMPLERATE,
|
|
threshold=VAD_THRESHOLD,
|
|
min_silence_duration_ms=MIN_SILENCE_DURATION_MS
|
|
)
|
|
|
|
if len(speech_timestamps) > 0:
|
|
speech_started = True
|
|
last_end = speech_timestamps[-1]['end']
|
|
buffer_len_samples = len(current_audio)
|
|
|
|
if (buffer_len_samples - last_end) > (SAMPLERATE * MIN_SILENCE_DURATION_MS / 1000) or buffer_len_samples > BUFFER_LIMIT:
|
|
|
|
# 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:
|
|
print(f"\n[{detected_lang.upper()} detected]: {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] + "..."
|
|
|
|
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
|
|
|
|
elif not speech_started and len(current_audio) > SAMPLERATE * 2:
|
|
audio_buffer = []
|
|
|
|
except KeyboardInterrupt:
|
|
print("\nStopped by user.")
|
|
except Exception as e:
|
|
print(f"\nError: {e}")
|
|
|
|
if __name__ == "__main__":
|
|
import multiprocessing
|
|
multiprocessing.freeze_support()
|
|
main()
|