102 lines
3.9 KiB
Python
102 lines
3.9 KiB
Python
import mlx_whisper
|
|
import numpy as np
|
|
import sounddevice as sd
|
|
import queue
|
|
import sys
|
|
import torch
|
|
from silero_vad import load_silero_vad, get_speech_timestamps
|
|
from mlx_lm import load, generate
|
|
|
|
# Parameters
|
|
WHISPER_MODEL = "mlx-community/whisper-small.en-mlx"
|
|
LLM_MODEL = "mlx-community/SmolLM2-135M-Instruct-4bit" # Tiny but fast LLM for translation
|
|
TARGET_LANG = "Spanish"
|
|
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 translate_text(model, tokenizer, text, target_lang):
|
|
prompt = f"Translate the following English text to {target_lang}. Only provide the translation, no extra text.\n\nEnglish: {text}\n\n{target_lang}:"
|
|
messages = [{"role": "user", "content": prompt}]
|
|
prompt_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
|
|
|
|
response = generate(model, tokenizer, prompt=prompt_text, max_tokens=100, verbose=False)
|
|
return response.strip()
|
|
|
|
def main():
|
|
print(f"Loading Whisper model '{WHISPER_MODEL}'...")
|
|
|
|
print(f"Loading LLM for translation '{LLM_MODEL}'...")
|
|
llm_model, llm_tokenizer = load(LLM_MODEL)
|
|
|
|
print("Loading Silero VAD model...")
|
|
vad_model = load_silero_vad()
|
|
|
|
print("Models loaded.")
|
|
|
|
print(f"\nStarting live transcription & translation to {TARGET_LANG}... (Press Ctrl+C to stop)")
|
|
|
|
audio_buffer = []
|
|
speech_started = False
|
|
|
|
try:
|
|
with sd.InputStream(samplerate=SAMPLERATE, channels=CHANNELS, callback=callback, blocksize=BLOCK_SIZE):
|
|
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 (English)
|
|
result = mlx_whisper.transcribe(current_audio, path_or_hf_repo=WHISPER_MODEL)
|
|
original_text = result['text'].strip()
|
|
|
|
if original_text:
|
|
# 2. Translate using LLM
|
|
translated_text = translate_text(llm_model, llm_tokenizer, original_text, TARGET_LANG)
|
|
|
|
print(f"\nEN: {original_text}")
|
|
print(f"{TARGET_LANG[: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__":
|
|
main()
|