whisper-translation/main_v2.py

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, run_llm_prompt_test
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", "qwen3.5:0.8b"))
    parser.add_argument("--post-correct-ollama-url", type=str, default=defaults.get("post_correct_ollama_url", "http://127.0.0.1:11434/api/generate"))
    parser.add_argument("--post-correct-llm-timeout", type=float, default=defaults.get("post_correct_llm_timeout", 8.0))
    parser.add_argument("--post-correct-keep-alive", type=str, default=defaults.get("post_correct_keep_alive", "30m"))
    parser.add_argument("--post-correct-warmup-timeout", type=float, default=defaults.get("post_correct_warmup_timeout", 20.0))
    parser.add_argument("--post-correct-min-overlap", type=float, default=defaults.get("post_correct_min_overlap", 0.45))
    parser.add_argument("--post-correct-debug", action="store_true", default=defaults.get("post_correct_debug", False))
    parser.add_argument("--freeflow-polish", action="store_true", default=defaults.get("freeflow_polish", True), help="Use Freeflow-style deterministic punctuation/filler polishing before optional LLM cleanup.")
    parser.add_argument("--no-freeflow-polish", action="store_false", dest="freeflow_polish", help="Use the older local regex cleanup instead of Freeflow-style deterministic polishing.")
    parser.add_argument("--post-correct-skip-clean", action="store_true", default=defaults.get("post_correct_skip_clean", True), help="Skip the line LLM when deterministic polish already looks clean.")
    parser.add_argument("--no-post-correct-skip-clean", action="store_false", dest="post_correct_skip_clean", help="Always call the line LLM when --post-correct-llm is enabled.")
    parser.add_argument("--post-correct-prompt-style", choices=["freeflow", "qwen", "legacy"], default=defaults.get("post_correct_prompt_style", "freeflow"), help="Prompt style for line LLM cleanup.")
    parser.add_argument("--llm-paragraph", action="store_true", default=defaults.get("llm_paragraph", False))
    parser.add_argument("--llm-paragraph-temperature", type=float, default=defaults.get("llm_paragraph_temperature", 0.2))
    parser.add_argument("--llm-paragraph-top-p", type=float, default=defaults.get("llm_paragraph_top_p", 0.8))
    parser.add_argument("--llm-paragraph-repeat-penalty", type=float, default=defaults.get("llm_paragraph_repeat_penalty", 1.0))
    parser.add_argument("--llm-request-log-path", type=str, default=defaults.get("llm_request_log_path", "logs/llm_requests.jsonl"))
    parser.add_argument("--llm-test-line", type=str, default=None, help="Send a single line-correction prompt without starting transcription.")
    parser.add_argument("--llm-test-prev1", type=str, default="", help="Optional previous line 1 context for --llm-test-line.")
    parser.add_argument("--llm-test-prev2", type=str, default="", help="Optional previous line 2 context for --llm-test-line.")
    parser.add_argument("--llm-test-segments", type=str, default=None, help="Send a single paragraph-refinement prompt without starting transcription.")
    parser.add_argument("--llm-test-context", type=str, default="", help="Optional existing paragraph context for --llm-test-segments.")
    parser.add_argument("--llm-test-from-log", type=int, default=None, help="Replay a logged LLM request by JSONL entry index. Use -1 for the most recent entry.")
    parser.add_argument("--llm-test-use-logged-model", action="store_true", default=False, help="When replaying from log, use the model stored in the log entry instead of --post-correct-model.")

    # 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

    if args.llm_test_line or args.llm_test_segments or args.llm_test_from_log is not None:
        run_llm_prompt_test(args)
        return

    # 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()