Adolfo Reyna
127 lines
5.8 KiB
Python
127 lines
5.8 KiB
Python
import multiprocessing
|
|
import argparse
|
|
import json
|
|
import os
|
|
import sys
|
|
from dotenv import load_dotenv
|
|
|
|
# Load environment variables from .env if present
|
|
load_dotenv()
|
|
|
|
# Light imports (heavy ones moved inside run functions)
|
|
from engine_transcribe import run_transcription, list_audio_devices
|
|
from engine_llm import run_llm_processor
|
|
from engine_translate import run_translation
|
|
from engine_distribute import run_distribution
|
|
|
|
def parse_temperature_fallback(value):
|
|
try:
|
|
return tuple(float(x.strip()) for x in value.split(",") if x.strip())
|
|
except:
|
|
return (0.0, 0.2, 0.4, 0.6, 0.8, 1.0)
|
|
|
|
def main():
|
|
config_path = "config.json"
|
|
defaults = {}
|
|
if os.path.exists(config_path):
|
|
with open(config_path, "r") as f:
|
|
defaults = json.load(f)
|
|
|
|
parser = argparse.ArgumentParser(description="Multi-process Transcription and Translation.")
|
|
|
|
# Transcribe Args
|
|
parser.add_argument("-l", "--list-devices", action="store_true", help="Show available audio devices and exit")
|
|
parser.add_argument("--model", type=str, default=defaults.get("model", "mlx-community/whisper-base-mlx"))
|
|
parser.add_argument("--device", type=int, default=defaults.get("device"))
|
|
parser.add_argument("--lang", type=str, default=defaults.get("lang"))
|
|
parser.add_argument("--silence", type=int, default=defaults.get("silence", 1000))
|
|
parser.add_argument("--max-buffer", type=int, default=defaults.get("max_buffer", 20))
|
|
parser.add_argument("--channels", type=int, default=defaults.get("channels", 1))
|
|
parser.add_argument("--filter-lang", action="store_true", default=defaults.get("filter_lang", False))
|
|
parser.add_argument("-q", "--quantize", action="store_true", default=defaults.get("quantize", False))
|
|
parser.add_argument("-s", "--stream", action="store_true", default=defaults.get("stream", False))
|
|
parser.add_argument("-c", "--context", action="store_true", default=defaults.get("context", False))
|
|
parser.add_argument("--speaker-diarization", action="store_true", default=defaults.get("speaker_diarization", False))
|
|
parser.add_argument("--temperature-fallback", type=parse_temperature_fallback, default=tuple(defaults.get("temperature_fallback", [0.0, 0.2, 0.4, 0.6, 0.8, 1.0])))
|
|
parser.add_argument("--logprob-threshold", type=float, default=defaults.get("logprob_threshold", -0.8))
|
|
parser.add_argument("--compression-threshold", type=float, default=defaults.get("compression_threshold", 2.2))
|
|
|
|
# LLM Args
|
|
parser.add_argument("--post-correct", action="store_true", default=defaults.get("post_correct", False))
|
|
parser.add_argument("--post-correct-llm", action="store_true", default=defaults.get("post_correct_llm", False))
|
|
parser.add_argument("--post-correct-model", type=str, default=defaults.get("post_correct_model", "qwen:2b"))
|
|
parser.add_argument("--llm-paragraph", action="store_true", default=defaults.get("llm_paragraph", False))
|
|
|
|
# Translate Args
|
|
parser.add_argument("-es", action="store_true", default=defaults.get("es", False))
|
|
parser.add_argument("-fr", action="store_true", default=defaults.get("fr", False))
|
|
parser.add_argument("-ar", action="store_true", default=defaults.get("ar", False))
|
|
parser.add_argument("-en", action="store_true", default=defaults.get("en", False))
|
|
parser.add_argument("--only-translate-llm", action="store_true", default=defaults.get("only_translate_llm", False), help="Only translate when the LLM produces a refined paragraph.")
|
|
parser.add_argument("--mt-max-chars", type=int, default=defaults.get("mt_max_chars", 250))
|
|
parser.add_argument("--mt-max-new-tokens", type=int, default=defaults.get("mt_max_new_tokens", 150))
|
|
parser.add_argument("--mt-num-beams", type=int, default=defaults.get("mt_num_beams", 4))
|
|
|
|
# Distribute Args
|
|
parser.add_argument("-i", "--ingest", action="store_true", default=defaults.get("ingest", False))
|
|
|
|
args = parser.parse_args()
|
|
|
|
# Auto-enable required flags
|
|
if args.post_correct_llm and not args.post_correct:
|
|
args.post_correct = True
|
|
if args.only_translate_llm and not args.llm_paragraph:
|
|
print("[Main] Enabling --llm-paragraph because --only-translate-llm was requested.")
|
|
args.llm_paragraph = True
|
|
|
|
# Handle device listing
|
|
if args.list_devices:
|
|
list_audio_devices()
|
|
return
|
|
|
|
# Handle interactive device selection
|
|
if args.device is None:
|
|
list_audio_devices()
|
|
try:
|
|
val = input("\nSelect input device index (or press Enter for default): ").strip()
|
|
if val:
|
|
args.device = int(val)
|
|
except EOFError:
|
|
pass # Non-interactive environment
|
|
except ValueError:
|
|
print("[Main] Invalid index, using default.")
|
|
|
|
# Queues
|
|
q_trans_to_llm = multiprocessing.Queue()
|
|
q_llm_to_tl = multiprocessing.Queue()
|
|
q_tl_to_dist = multiprocessing.Queue()
|
|
|
|
# Processes
|
|
p_transcribe = multiprocessing.Process(target=run_transcription, args=(q_trans_to_llm, args))
|
|
p_llm = multiprocessing.Process(target=run_llm_processor, args=(q_trans_to_llm, q_llm_to_tl, args))
|
|
p_translate = multiprocessing.Process(target=run_translation, args=(q_llm_to_tl, q_tl_to_dist, args))
|
|
p_distribute = multiprocessing.Process(target=run_distribution, args=(q_tl_to_dist, args))
|
|
|
|
print("[Main] Starting processes...")
|
|
p_transcribe.start()
|
|
p_llm.start()
|
|
p_translate.start()
|
|
p_distribute.start()
|
|
|
|
try:
|
|
p_transcribe.join()
|
|
p_llm.join()
|
|
p_translate.join()
|
|
p_distribute.join()
|
|
except KeyboardInterrupt:
|
|
print("\n[Main] Stopping processes...")
|
|
p_transcribe.terminate()
|
|
p_llm.terminate()
|
|
p_translate.terminate()
|
|
p_distribute.terminate()
|
|
sys.exit(0)
|
|
|
|
if __name__ == "__main__":
|
|
multiprocessing.freeze_support()
|
|
main()
|