feat: add command-line arguments for lang selection, ingest toggle, and device management; add README and requirements
This commit is contained in:
Adolfo Reyna
+54
-21
@@ -3,6 +3,7 @@ import time
|
||||
import requests
|
||||
import threading
|
||||
import json
|
||||
import argparse
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
# Comprehensive workaround for missing _lzma in some Python builds
|
||||
@@ -82,17 +83,37 @@ def ingest_worker():
|
||||
ingest_queue.task_done()
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Live transcription and translation with Whisper.")
|
||||
parser.add_argument("-es", action="store_true", help="Enable Spanish translation")
|
||||
parser.add_argument("-en", action="store_true", help="Enable English (detected or bridged)")
|
||||
parser.add_argument("-ar", action="store_true", help="Enable Arabic translation")
|
||||
parser.add_argument("-fr", action="store_true", help="Enable French translation")
|
||||
parser.add_argument("-i", "--ingest", action="store_true", help="Enable data transmission to server")
|
||||
parser.add_argument("-l", "--list-devices", action="store_true", help="Show available audio devices and exit")
|
||||
parser.add_argument("-d", "--device", type=int, help="Input device index")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.list_devices:
|
||||
print("\nAvailable Audio Devices:")
|
||||
print(sd.query_devices())
|
||||
return
|
||||
|
||||
device = "mps" if torch.backends.mps.is_available() else "cpu"
|
||||
print(f"Using device: {device}")
|
||||
|
||||
# Start ingest thread
|
||||
threading.Thread(target=ingest_worker, daemon=True).start()
|
||||
# Only start ingest thread if enabled
|
||||
if args.ingest:
|
||||
threading.Thread(target=ingest_worker, daemon=True).start()
|
||||
|
||||
# 1. Load models
|
||||
print(f"Loading Multilingual Whisper model '{WHISPER_MODEL}'...")
|
||||
|
||||
# Filter translation models to only those enabled by args
|
||||
active_target_langs = {k: v for k, v in TARGET_LANGS.items() if getattr(args, k, False)}
|
||||
|
||||
translation_engines = {}
|
||||
for lang_key, model_id in TARGET_LANGS.items():
|
||||
for lang_key, model_id in active_target_langs.items():
|
||||
print(f"Loading {lang_key} translation model ({model_id})...")
|
||||
tokenizer = MarianTokenizer.from_pretrained(model_id)
|
||||
model = MarianMTModel.from_pretrained(model_id).to(device)
|
||||
@@ -103,17 +124,19 @@ def main():
|
||||
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
|
||||
if args.device is not None:
|
||||
device_index = args.device
|
||||
else:
|
||||
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 live transcription & server ingest... (Press Ctrl+C to stop)")
|
||||
print(f"\nStarting live transcription{' & server ingest' if args.ingest else ''}... (Press Ctrl+C to stop)")
|
||||
|
||||
audio_buffer = []
|
||||
speech_started = False
|
||||
@@ -154,36 +177,46 @@ def main():
|
||||
|
||||
# Prepare payload
|
||||
payload = {"original": original_text}
|
||||
# Rule 3: include source language key
|
||||
if detected_lang in TARGET_LANGS or detected_lang == "en":
|
||||
|
||||
# Include detected language if requested or if it's the bridge
|
||||
if (detected_lang in active_target_langs) or (detected_lang == "en" and args.en):
|
||||
payload[detected_lang] = original_text
|
||||
|
||||
# 2. Bridge to English if not already English
|
||||
if detected_lang != "en":
|
||||
bridge_result = mlx_whisper.transcribe(current_audio, path_or_hf_repo=WHISPER_MODEL, task="translate")
|
||||
english_text = bridge_result['text'].strip()
|
||||
payload["en"] = english_text
|
||||
if args.en:
|
||||
payload["en"] = english_text
|
||||
print(f"[EN]: {english_text}")
|
||||
else:
|
||||
english_text = original_text
|
||||
if args.en:
|
||||
payload["en"] = english_text
|
||||
|
||||
# 3. Translate from English to other languages
|
||||
if english_text:
|
||||
if english_text and translation_engines:
|
||||
# Limit input length
|
||||
clean_en = english_text[:247] + "..." if len(english_text) > 250 else english_text
|
||||
|
||||
for lang_key, (model, tokenizer) in translation_engines.items():
|
||||
# Skip if we already filled this (e.g. detected lang was 'es')
|
||||
if lang_key in payload: continue
|
||||
if lang_key in payload:
|
||||
if lang_key != detected_lang: # Already printed original
|
||||
print(f"[{lang_key.upper()}]: {payload[lang_key]}")
|
||||
continue
|
||||
|
||||
inputs = tokenizer(clean_en, 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)
|
||||
payload[lang_key] = translated_text
|
||||
print(f"[{lang_key.upper()}]: {translated_text}")
|
||||
|
||||
# Queue for background ingestion
|
||||
ingest_queue.put(payload)
|
||||
print(f"Sent to ingest: {list(payload.keys())}")
|
||||
# Queue for background ingestion if enabled
|
||||
if args.ingest:
|
||||
ingest_queue.put(payload)
|
||||
# print(f"Sent to ingest: {list(payload.keys())}")
|
||||
|
||||
audio_buffer = []
|
||||
speech_started = False
|
||||
|
||||
Reference in New Issue
Block a user