live_transcript.py

#!/usr/bin/env python3
"""
Live desktop-audio transcription (Ubuntu / PipeWire/PulseAudio).

All artifacts are stored relative to THIS FILE:
  ./models/                 whisper models
  ./chunks/                 generated WAV audio chunks
  ./transcript_live.txt     continuously appended transcript
  ./asr_consumer.log        operational logs

Prereqs:
  - pactl null sink setup (asr_sink + optional loopback)
  - ffmpeg installed
  - whisper.cpp built (whisper-cli binary available)

Run:
  python3 live_transcript.py

Stop:
  Ctrl+C
"""

import os
import signal
import shutil
import subprocess
import time
import atexit
import ctypes
import re
import argparse
from datetime import datetime, timedelta
from pathlib import Path

# ---------- PATH SETUP (RELATIVE TO SCRIPT) ----------

BASE_DIR = Path(__file__).resolve().parent
MODELS_DIR = BASE_DIR / "models"
CHUNKS_DIR = BASE_DIR / "chunks"

TRANSCRIPT_PATH = BASE_DIR / "transcript_live.txt"
LOG_PATH = BASE_DIR / "asr_consumer.log"

# ---------- CONFIG ----------

PULSE_SOURCE = "asr_sink.monitor"
SEGMENT_SECONDS = 6
KEEP_LATEST_CHUNKS = 20
MAX_BACKLOG_CHUNKS = 4
TARGET_BACKLOG_CHUNKS = 1
DEFAULT_LANGUAGE_MODE = os.getenv("ASR_LANGUAGE_MODE", "auto")
DEFAULT_LANGUAGE = os.getenv("ASR_LANGUAGE", "en")
DEFAULT_LANGUAGES = os.getenv("ASR_LANGUAGES", "en,pt")
DEFAULT_TRANSLATION_MODE = os.getenv("ASR_TRANSLATION_MODE", "original")
DEFAULT_ACCURACY_MODE = os.getenv("ASR_ACCURACY_MODE", "balanced")
ALLOW_ENGLISH_FALLBACK = os.getenv("ASR_ALLOW_ENGLISH_FALLBACK", "0") == "1"
AUTO_DOWNLOAD_MULTILINGUAL_MODEL = os.getenv("ASR_AUTO_DOWNLOAD_MODEL", "1") == "1"

FFMPEG_BIN = "ffmpeg"
MODEL_NAME = "ggml-base.en.bin"
MODEL_PATH = MODELS_DIR / MODEL_NAME

STABLE_CHECK_SLEEP = 0.30
MIN_NONZERO_BYTES = 1
POLL_INTERVAL_SECONDS = 0.10
TIMESTAMP_RE = re.compile(
    r"^\[(\d+):(\d+):(\d+(?:\.\d+)?)\s*-->\s*(\d+):(\d+):(\d+(?:\.\d+)?)\]\s*(.*)$"
)

# ---------- UTILS ----------

def log(msg: str) -> None:
    ts = time.strftime("%Y-%m-%d %H:%M:%S")
    with LOG_PATH.open("a", encoding="utf-8") as f:
        f.write(f"[{ts}] {msg}\n")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Live desktop audio transcription")
    parser.add_argument(
        "--language-mode",
        choices=["auto", "single", "multi-pass"],
        default=DEFAULT_LANGUAGE_MODE,
        help="auto: language auto-detect (best for mixed EN/PT), single: force one language, multi-pass: run one pass per language and pick best",
    )
    parser.add_argument(
        "--language",
        default=DEFAULT_LANGUAGE,
        help="language code used in single mode (example: en, pt, pt-br)",
    )
    parser.add_argument(
        "--languages",
        default=DEFAULT_LANGUAGES,
        help="comma-separated language codes for multi-pass mode (example: en,pt)",
    )
    parser.add_argument(
        "--translation-mode",
        choices=["original", "english"],
        default=DEFAULT_TRANSLATION_MODE,
        help="original: keep detected language text, english: translate output to English",
    )
    parser.add_argument(
        "--accuracy-mode",
        choices=["fast", "balanced", "high"],
        default=DEFAULT_ACCURACY_MODE,
        help="fast: lowest latency, balanced: better quality, high: best quality (slower)",
    )
    return parser.parse_args()


def normalize_language_code(lang: str) -> str:
    normalized = lang.strip().lower().replace("_", "-")
    aliases = {
        "pt-br": "pt",
        "pt-pt": "pt",
        "en-us": "en",
        "en-gb": "en",
    }
    return aliases.get(normalized, normalized)


def parse_languages(languages_csv: str) -> list[str]:
    langs = [normalize_language_code(lang) for lang in languages_csv.split(",")]
    langs = [lang for lang in langs if lang]
    return langs


def resolve_language_config(args: argparse.Namespace) -> tuple[str, list[str]]:
    mode = args.language_mode.strip().lower()
    single_lang = normalize_language_code(args.language)
    multi_langs = parse_languages(args.languages)

    if mode == "single":
        return "single", [single_lang]
    if mode == "multi-pass":
        if not multi_langs:
            raise RuntimeError("multi-pass mode requires at least one language in --languages")
        return "multi-pass", multi_langs
    return "auto", ["auto"]


def resolve_decode_params(accuracy_mode: str) -> tuple[str, str]:
    if accuracy_mode == "high":
        return "5", "5"
    if accuracy_mode == "balanced":
        return "3", "3"
    return "1", "1"


def try_upgrade_to_multilingual_model(
    model_path: Path, language_mode: str, languages: list[str]
) -> Path:
    needs_multilingual = language_mode != "single" or (languages and languages[0] != "en")
    if not needs_multilingual:
        return model_path

    if not model_path.name.endswith(".en.bin"):
        return model_path

    multilingual_candidates: list[Path] = []
    de_english_name = model_path.name.replace(".en.bin", ".bin")
    multilingual_candidates.extend(
        [
            MODELS_DIR / de_english_name,
            Path.home() / "whisper.cpp" / "models" / de_english_name,
            MODELS_DIR / "ggml-base.bin",
            Path.home() / "whisper.cpp" / "models" / "ggml-base.bin",
        ]
    )

    for candidate in multilingual_candidates:
        if not candidate.exists():
            continue

        target = candidate if candidate.parent == MODELS_DIR else MODELS_DIR / candidate.name
        if candidate != target and not target.exists():
            shutil.copy2(candidate, target)
            log(f"Copied multilingual model to {target}")
        log(f"Using multilingual model for language mode {language_mode}: {target}")
        return target

    if AUTO_DOWNLOAD_MULTILINGUAL_MODEL:
        downloader = BASE_DIR / "whisper.cpp" / "models" / "download-ggml-model.sh"
        if downloader.exists() and os.access(downloader, os.X_OK):
            log("Multilingual model missing. Attempting automatic download: base")
            proc = subprocess.run(
                [str(downloader), "base"],
                cwd=str(BASE_DIR / "whisper.cpp"),
                capture_output=True,
                text=True,
            )
            with LOG_PATH.open("a", encoding="utf-8") as logf:
                if proc.stdout:
                    logf.write(proc.stdout)
                if proc.stderr:
                    logf.write(proc.stderr)
            if proc.returncode == 0:
                downloaded = BASE_DIR / "whisper.cpp" / "models" / "ggml-base.bin"
                if downloaded.exists():
                    target = MODELS_DIR / downloaded.name
                    if not target.exists():
                        shutil.copy2(downloaded, target)
                        log(f"Copied multilingual model to {target}")
                    log(f"Using multilingual model for language mode {language_mode}: {target}")
                    return target
                log("Automatic download completed but ggml-base.bin was not found.")

    msg = (
        "Multilingual mode requested, but only English-only model is available. "
        "Download a multilingual model like ggml-base.bin (whisper.cpp/models) "
        "or set WHISPER_MODEL to a multilingual .bin file."
    )
    if ALLOW_ENGLISH_FALLBACK:
        log(f"Warning: {msg} Falling back to English-only model because ASR_ALLOW_ENGLISH_FALLBACK=1.")
        return model_path
    raise RuntimeError(msg)


def _set_pdeathsig() -> None:
    # Ensure ffmpeg gets SIGTERM automatically if this Python process dies.
    libc = ctypes.CDLL("libc.so.6")
    PR_SET_PDEATHSIG = 1
    libc.prctl(PR_SET_PDEATHSIG, signal.SIGTERM)


def resolve_whisper_bin() -> Path:
    env_bin = os.getenv("WHISPER_BIN")
    candidates = []
    if env_bin:
        candidates.append(Path(env_bin).expanduser())
    candidates.extend(
        [
            BASE_DIR / "whisper.cpp" / "build" / "bin" / "whisper-cli",
            Path.home() / "whisper.cpp" / "build" / "bin" / "whisper-cli",
        ]
    )

    for candidate in candidates:
        if candidate.exists() and os.access(candidate, os.X_OK):
            return candidate

    raise RuntimeError(
        "whisper-cli not found. Set WHISPER_BIN or install whisper.cpp."
    )


def resolve_model_path() -> Path:
    MODELS_DIR.mkdir(exist_ok=True)
    env_model = os.getenv("WHISPER_MODEL")
    if env_model:
        src = Path(env_model).expanduser()
        if not src.exists():
            raise RuntimeError(f"WHISPER_MODEL does not exist: {src}")
        if src != MODEL_PATH:
            shutil.copy2(src, MODEL_PATH)
            log(f"Copied model to {MODEL_PATH}")
        return MODEL_PATH

    if MODEL_PATH.exists():
        return MODEL_PATH

    fallback_model = Path.home() / "whisper.cpp" / "models" / MODEL_NAME
    if fallback_model.exists():
        shutil.copy2(fallback_model, MODEL_PATH)
        log(f"Copied model to {MODEL_PATH}")
        return MODEL_PATH

    raise RuntimeError(
        f"Model not found at {MODEL_PATH}. Put {MODEL_NAME} in ./models or set WHISPER_MODEL."
    )


def build_whisper_env(whisper_bin: Path) -> dict[str, str]:
    env = os.environ.copy()
    build_dir = whisper_bin.parent.parent
    lib_dirs = []
    for lib_dir in [build_dir / "src", build_dir / "ggml" / "src"]:
        if lib_dir.exists():
            lib_dirs.append(str(lib_dir))

    if lib_dirs:
        current = env.get("LD_LIBRARY_PATH", "")
        env["LD_LIBRARY_PATH"] = ":".join(lib_dirs + ([current] if current else []))

    return env


def check_prereqs() -> tuple[Path, Path, dict[str, str]]:
    if shutil.which(FFMPEG_BIN) is None:
        raise RuntimeError("ffmpeg not found in PATH.")

    whisper_bin = resolve_whisper_bin()
    model_path = resolve_model_path()
    whisper_env = build_whisper_env(whisper_bin)

    probe = subprocess.run(
        [str(whisper_bin), "-h"],
        stdout=subprocess.DEVNULL,
        stderr=subprocess.DEVNULL,
        env=whisper_env,
    )
    if probe.returncode != 0:
        raise RuntimeError(
            "whisper-cli could not start. Check shared libraries (LD_LIBRARY_PATH)."
        )

    return whisper_bin, model_path, whisper_env


def wait_for_file_stable(path: Path) -> None:
    while not path.exists():
        time.sleep(0.2)

    while True:
        s1 = path.stat().st_size
        time.sleep(STABLE_CHECK_SLEEP)
        s2 = path.stat().st_size
        if s1 == s2 and s2 >= MIN_NONZERO_BYTES:
            return


# ---------- FFMPEG PRODUCER ----------

def start_ffmpeg_chunker() -> tuple[subprocess.Popen, Path]:
    CHUNKS_DIR.mkdir(exist_ok=True)
    run_chunks_dir = CHUNKS_DIR / f"run_{int(time.time())}"
    run_chunks_dir.mkdir(exist_ok=True)

    cmd = [
        FFMPEG_BIN,
        "-nostdin",
        "-f", "pulse",
        "-i", PULSE_SOURCE,
        "-ac", "1",
        "-ar", "16000",
        "-c:a", "pcm_s16le",
        "-f", "segment",
        "-segment_time", str(SEGMENT_SECONDS),
        "-reset_timestamps", "1",
        str(run_chunks_dir / "chunk_%06d.wav"),
    ]

    log(f"Starting ffmpeg chunker in {run_chunks_dir}")
    proc = subprocess.Popen(
        cmd,
        stdout=subprocess.DEVNULL,
        stderr=subprocess.PIPE,
        text=True,
        preexec_fn=_set_pdeathsig,
    )
    return proc, run_chunks_dir


def stop_ffmpeg_proc(ffmpeg_proc: subprocess.Popen | None) -> None:
    if ffmpeg_proc is None:
        return
    if ffmpeg_proc.poll() is None:
        ffmpeg_proc.send_signal(signal.SIGINT)
        try:
            ffmpeg_proc.wait(timeout=5)
        except subprocess.TimeoutExpired:
            ffmpeg_proc.kill()


def cleanup_stale_chunk_writers() -> None:
    # Kill only ffmpeg processes writing chunks into this project's chunks dir.
    base = str(CHUNKS_DIR)
    result = subprocess.run(
        ["pgrep", "-af", "ffmpeg"],
        capture_output=True,
        text=True,
    )
    if result.returncode not in (0, 1):
        return

    for line in result.stdout.splitlines():
        if base not in line or "chunk_%06d.wav" not in line:
            continue
        try:
            pid = int(line.split(maxsplit=1)[0])
        except (ValueError, IndexError):
            continue
        if pid == os.getpid():
            continue
        try:
            os.kill(pid, signal.SIGTERM)
            log(f"Killed stale ffmpeg chunk writer pid={pid}")
        except ProcessLookupError:
            pass


# ---------- WHISPER CONSUMER ----------

def _hms_to_seconds(hours: str, minutes: str, seconds: str) -> float:
    return int(hours) * 3600 + int(minutes) * 60 + float(seconds)


def _fmt_absolute(ts: datetime) -> str:
    millis = ts.microsecond // 1000
    return f"{ts:%Y-%m-%d %H:%M:%S}.{millis:03d}"


def convert_to_absolute_timestamps(raw_text: str, chunk_start: datetime) -> str:
    lines: list[str] = []
    for raw_line in raw_text.splitlines():
        line = raw_line.strip()
        if not line:
            continue

        match = TIMESTAMP_RE.match(line)
        if not match:
            lines.append(line)
            continue

        start_s = _hms_to_seconds(match.group(1), match.group(2), match.group(3))
        text = match.group(7).strip()

        abs_start = chunk_start + timedelta(seconds=start_s)
        lines.append(f"[{_fmt_absolute(abs_start)}] {text}")

    if not lines:
        return ""
    return "\n".join(lines) + "\n"


def chunk_index(chunk: Path) -> int | None:
    match = re.match(r"chunk_(\d+)\.wav$", chunk.name)
    if not match:
        return None
    return int(match.group(1))


def get_ready_chunk_indexes(run_chunks_dir: Path) -> list[int]:
    indexes: list[int] = []
    for chunk in run_chunks_dir.glob("chunk_*.wav"):
        idx = chunk_index(chunk)
        if idx is None:
            continue
        if not chunk.exists() or chunk.stat().st_size <= 0:
            continue
        indexes.append(idx)
    indexes.sort()
    return indexes


def prune_old_chunks(run_chunks_dir: Path, last_processed_idx: int) -> None:
    keep_from = max(0, last_processed_idx - KEEP_LATEST_CHUNKS + 1)
    for chunk in run_chunks_dir.glob("chunk_*.wav"):
        idx = chunk_index(chunk)
        if idx is None or idx >= keep_from:
            continue
        try:
            chunk.unlink()
        except FileNotFoundError:
            pass


def score_transcript_output(converted: str) -> tuple[int, int]:
    speech_lines = 0
    speech_chars = 0
    for line in converted.splitlines():
        if not line or "[BLANK_AUDIO]" in line:
            continue
        speech_lines += 1
        speech_chars += len(line)
    return speech_lines, speech_chars


def run_whisper(
    chunk: Path,
    whisper_bin: Path,
    model_path: Path,
    whisper_env: dict[str, str],
    language: str,
    translate_to_english: bool,
    best_of: str,
    beam_size: str,
) -> subprocess.CompletedProcess[str]:
    cmd = [
        str(whisper_bin),
        "-m", str(model_path),
        "-l", language,
        "-bo", best_of,
        "-bs", beam_size,
        "-t", str(max(1, os.cpu_count() or 1)),
        "-f", str(chunk),
    ]
    if translate_to_english:
        cmd.insert(6, "-tr")
    return subprocess.run(cmd, capture_output=True, text=True, env=whisper_env)


def transcribe_chunk(
    chunk: Path,
    whisper_bin: Path,
    model_path: Path,
    whisper_env: dict[str, str],
    chunk_start: datetime,
    language_mode: str,
    languages: list[str],
    translate_to_english: bool,
    best_of: str,
    beam_size: str,
) -> None:
    if language_mode == "multi-pass":
        best_converted = ""
        best_score = (-1, -1)
        best_lang = ""
        for lang in languages:
            proc = run_whisper(
                chunk,
                whisper_bin,
                model_path,
                whisper_env,
                lang,
                translate_to_english,
                best_of,
                beam_size,
            )
            with LOG_PATH.open("a", encoding="utf-8") as logf:
                if proc.stderr:
                    logf.write(proc.stderr)
            if proc.returncode != 0:
                log(f"whisper-cli failed for {chunk.name} lang={lang} (code={proc.returncode})")
                continue

            converted = convert_to_absolute_timestamps(proc.stdout, chunk_start)
            score = score_transcript_output(converted)
            if score > best_score:
                best_score = score
                best_converted = converted
                best_lang = lang

        if best_lang:
            log(f"Selected language {best_lang} for {chunk.name} score={best_score}")
            converted = best_converted
        else:
            return
    else:
        lang = languages[0]
        proc = run_whisper(
            chunk,
            whisper_bin,
            model_path,
            whisper_env,
            lang,
            translate_to_english,
            best_of,
            beam_size,
        )
        with LOG_PATH.open("a", encoding="utf-8") as logf:
            if proc.stderr:
                logf.write(proc.stderr)
        if proc.returncode != 0:
            log(f"whisper-cli failed for {chunk.name} lang={lang} (code={proc.returncode})")
            return
        converted = convert_to_absolute_timestamps(proc.stdout, chunk_start)

    if converted:
        with TRANSCRIPT_PATH.open("a", encoding="utf-8") as out:
            out.write(converted)


# ---------- MAIN ----------

def main() -> None:
    MODELS_DIR.mkdir(exist_ok=True)
    CHUNKS_DIR.mkdir(exist_ok=True)
    TRANSCRIPT_PATH.touch(exist_ok=True)
    LOG_PATH.touch(exist_ok=True)

    args = parse_args()
    language_mode, languages = resolve_language_config(args)
    translate_to_english = args.translation_mode == "english"
    best_of, beam_size = resolve_decode_params(args.accuracy_mode)

    whisper_bin, model_path, whisper_env = check_prereqs()
    model_path = try_upgrade_to_multilingual_model(model_path, language_mode, languages)
    log(f"Using whisper binary: {whisper_bin}")
    log(f"Using model file: {model_path}")
    log(f"Language mode: {language_mode} languages={','.join(languages)}")
    log(f"Translation mode: {'english' if translate_to_english else 'original'}")
    log(f"Accuracy mode: {args.accuracy_mode} (best_of={best_of}, beam_size={beam_size})")

    cleanup_stale_chunk_writers()
    ffmpeg_proc, run_chunks_dir = start_ffmpeg_chunker()
    run_started_at = datetime.now()

    def shutdown(*_):
        log("Shutdown requested")
        stop_ffmpeg_proc(ffmpeg_proc)
        raise SystemExit(0)

    atexit.register(lambda: stop_ffmpeg_proc(ffmpeg_proc))
    signal.signal(signal.SIGINT, shutdown)
    signal.signal(signal.SIGTERM, shutdown)
    signal.signal(signal.SIGHUP, shutdown)

    i = 0
    while True:
        if ffmpeg_proc.poll() is not None:
            err = ffmpeg_proc.stderr.read() if ffmpeg_proc.stderr else ""
            log(f"ffmpeg exited unexpectedly:\n{err}")
            raise SystemExit(1)

        ready = get_ready_chunk_indexes(run_chunks_dir)
        if not ready:
            time.sleep(POLL_INTERVAL_SECONDS)
            continue

        latest_ready = ready[-1]
        backlog = latest_ready - i
        if backlog > MAX_BACKLOG_CHUNKS:
            new_i = max(0, latest_ready - TARGET_BACKLOG_CHUNKS)
            log(
                f"Backlog detected ({backlog} chunks). Skipping from chunk_{i:06d} to chunk_{new_i:06d}."
            )
            i = new_i

        chunk = run_chunks_dir / f"chunk_{i:06d}.wav"
        if not chunk.exists():
            time.sleep(POLL_INTERVAL_SECONDS)
            continue

        wait_for_file_stable(chunk)
        log(f"Transcribing {chunk.name} ({chunk.stat().st_size} bytes)")
        chunk_start = run_started_at + timedelta(seconds=i * SEGMENT_SECONDS)
        transcribe_chunk(
            chunk,
            whisper_bin,
            model_path,
            whisper_env,
            chunk_start,
            language_mode,
            languages,
            translate_to_english,
            best_of,
            beam_size,
        )
        prune_old_chunks(run_chunks_dir, i)
        i += 1


if __name__ == "__main__":
    main()