adolforeyna/whisper-translation
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Switch to mlx-lm for translation to avoid LZMA dependency error

Browse Source
This commit is contained in:
Adolfo Reyna
2026-02-26 21:09:34 -05:00
parent 035d6e9358
commit 78322cbc8d
1 changed files with 18 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files

33
transcribe.py
View File

@@ -5,12 +5,12 @@ import queue
import sys import sys
import torch import torch
from silero_vad import load_silero_vad, get_speech_timestamps from silero_vad import load_silero_vad, get_speech_timestamps
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline from mlx_lm import load, generate
# Parameters # Parameters
WHISPER_MODEL = "mlx-community/whisper-small.en-mlx" WHISPER_MODEL = "mlx-community/whisper-small.en-mlx"
TRANSLATE_MODEL = "facebook/nllb-200-distilled-600M" LLM_MODEL = "mlx-community/SmolLM2-135M-Instruct-4bit" # Tiny but fast LLM for translation
TARGET_LANG = "spa_Latn" # Spanish (Latin America) - change as needed TARGET_LANG = "Spanish"
CHANNELS = 1 CHANNELS = 1
SAMPLERATE = 16000 SAMPLERATE = 16000
BLOCK_SIZE = 512 BLOCK_SIZE = 512
@@ -25,19 +25,19 @@ def callback(indata, frames, time, status):
print(status, file=sys.stderr) print(status, file=sys.stderr)
audio_queue.put(indata.copy()) 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(): def main():
# Set device for translation model (using MPS for Mac M chips)
device = "mps" if torch.backends.mps.is_available() else "cpu"
print(f"Using device: {device}")
print(f"Loading Whisper model '{WHISPER_MODEL}'...") print(f"Loading Whisper model '{WHISPER_MODEL}'...")
# Whisper MLX runs on its own optimized path
print(f"Loading translation model '{TRANSLATE_MODEL}'...") print(f"Loading LLM for translation '{LLM_MODEL}'...")
tokenizer = AutoTokenizer.from_pretrained(TRANSLATE_MODEL) llm_model, llm_tokenizer = load(LLM_MODEL)
model = AutoModelForSeq2SeqLM.from_pretrained(TRANSLATE_MODEL).to(device)
translator = pipeline("translation", model=model, tokenizer=tokenizer,
src_lang="eng_Latn", tgt_lang=TARGET_LANG, device=device)
print("Loading Silero VAD model...") print("Loading Silero VAD model...")
vad_model = load_silero_vad() vad_model = load_silero_vad()
@@ -80,12 +80,11 @@ def main():
original_text = result['text'].strip() original_text = result['text'].strip()
if original_text: if original_text:
# 2. Translate # 2. Translate using LLM
translation = translator(original_text) translated_text = translate_text(llm_model, llm_tokenizer, original_text, TARGET_LANG)
translated_text = translation[0]['translation_text']
print(f"\nEN: {original_text}") print(f"\nEN: {original_text}")
print(f"ES: {translated_text}") print(f"{TARGET_LANG[:2].upper()}: {translated_text}")
audio_buffer = [] audio_buffer = []
speech_started = False speech_started = False