zyt/app/scripts/setup_local_transcription.py

"""Explicit installation/download/check CLI; never invoked by a live call.

Run with --help for examples. With no action flags this only prints help.
"""

from __future__ import annotations

import argparse
import importlib.util
import json
import os
import shutil
import subprocess
import sys
import wave
from datetime import UTC, datetime
from pathlib import Path
from unittest.mock import patch

APP_ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(APP_ROOT / "src"))

MODEL_REPOS = {
    size: f"Systran/faster-whisper-{size}"
    for size in ("tiny", "base", "small", "medium", "large-v3")
}
RUNTIME_REQUIREMENT = "faster-whisper==1.2.1"
TEXT_REQUIREMENT = "opencc-python-reimplemented==0.1.7"


def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Install or verify local CPU/int8 speech recognition before making calls.",
        epilog=(
            "Example: python scripts/setup_local_transcription.py "
            "--install-dependencies --download --model-size small --check. "
            "Offline check: python scripts/setup_local_transcription.py --check. "
            "A packaged app must ship the runtime; installing into another Python "
            "does not alter a frozen executable."
        ),
    )
    parser.add_argument(
        "--install-dependencies",
        action="store_true",
        help="Explicitly run pip in this Python environment.",
    )
    parser.add_argument(
        "--download",
        action="store_true",
        help="Explicitly download the selected official model from Hugging Face.",
    )
    parser.add_argument("--model-size", choices=tuple(MODEL_REPOS), default="small")
    parser.add_argument(
        "--model-path",
        type=Path,
        help="Override DOCTOR_ASR_MODEL_PATH or the app data model directory.",
    )
    parser.add_argument(
        "--revision", default="main", help="Official repository revision/commit for model download."
    )
    parser.add_argument(
        "--check",
        action="store_true",
        help="Load the local model and VAD on CPU with networking disabled.",
    )
    parser.add_argument(
        "--smoke-wav",
        type=Path,
        help="Optional non-sensitive mono 16 kHz PCM16 WAV, at most 12 seconds; recognized text is printed.",
    )
    args = parser.parse_args(argv)
    if not (args.install_dependencies or args.download or args.check or args.smoke_wav):
        parser.print_help()
    return args


def main(argv: list[str] | None = None) -> int:
    args = parse_args(argv)
    if not (args.install_dependencies or args.download or args.check or args.smoke_wav):
        return 0
    if args.install_dependencies:
        installer = [sys.executable, "-m", "pip", "install"]
        if importlib.util.find_spec("pip") is None:
            uv = shutil.which("uv")
            if not uv:
                print(
                    "This environment has no pip. Install with uv pip install --python <python> "
                    f"{RUNTIME_REQUIREMENT} {TEXT_REQUIREMENT}",
                    file=sys.stderr,
                )
                return 1
            installer = [uv, "pip", "install", "--python", sys.executable]
        subprocess.run(
            [*installer, RUNTIME_REQUIREMENT, TEXT_REQUIREMENT, "platformdirs>=4.3,<5"],
            check=True,
        )

    from doctor_workstation.core.transcript_text import simplified_transcript
    from doctor_workstation.video.local_transcription import (
        MAX_PCM_BYTES,
        SAMPLE_RATE,
        FasterWhisperEngine,
        LocalTranscriptionError,
        default_model_path,
        local_transcription_readiness,
    )

    path = args.model_path.expanduser().resolve() if args.model_path else default_model_path()
    try:
        if args.download:
            from huggingface_hub import snapshot_download

            path.mkdir(parents=True, exist_ok=True)
            repo = MODEL_REPOS[args.model_size]
            print(f"Downloading {repo} into {path}")
            snapshot_download(
                repo_id=repo,
                revision=args.revision,
                local_dir=str(path),
                allow_patterns=[
                    "config.json",
                    "preprocessor_config.json",
                    "model.bin",
                    "tokenizer.json",
                    "vocabulary.*",
                    "README.md",
                    "LICENSE*",
                ],
                token=False,
            )
            (path / "doctor-asr-install.json").write_text(
                json.dumps(
                    {
                        "repository": repo,
                        "requestedRevision": args.revision,
                        "runtimeRequirement": RUNTIME_REQUIREMENT,
                        "textRequirement": TEXT_REQUIREMENT,
                        "installedAt": datetime.now(UTC).isoformat(),
                    },
                    indent=2,
                )
                + "\n",
                encoding="utf-8",
            )
        if args.check or args.smoke_wav or args.download:
            # This affects only the explicit setup/check process. The runtime
            # already requires complete local files and local_files_only=True.
            os.environ["HF_HUB_OFFLINE"] = "1"
            os.environ["HF_HUB_DISABLE_TELEMETRY"] = "1"
            print(json.dumps(local_transcription_readiness(path), ensure_ascii=False, indent=2))
            # The hub may already have been imported during --download, so its
            # offline environment flag alone is not a sufficient verification.
            # Block sockets in this standalone check process as well.
            blocked = RuntimeError("Network access is disabled during local ASR verification")
            with (
                patch("socket.socket.connect", side_effect=blocked),
                patch("socket.socket.connect_ex", side_effect=blocked),
                patch("socket.getaddrinfo", side_effect=blocked),
            ):
                engine = FasterWhisperEngine(path)
                print("Local model and VAD loaded successfully (CPU/int8; network blocked).")
                if args.smoke_wav:
                    with wave.open(str(args.smoke_wav), "rb") as audio:
                        if (
                            audio.getnchannels(),
                            audio.getsampwidth(),
                            audio.getframerate(),
                            audio.getcomptype(),
                        ) != (1, 2, SAMPLE_RATE, "NONE"):
                            raise ValueError("Smoke WAV must be mono 16 kHz uncompressed PCM16.")
                        if not 0 < audio.getnframes() <= MAX_PCM_BYTES // 2:
                            raise ValueError(
                                "Smoke WAV must contain between one sample and 12 seconds."
                            )
                        pcm = audio.readframes(audio.getnframes())
                    print(f"Recognized text: {simplified_transcript(engine.transcribe(pcm))}")
        return 0
    except LocalTranscriptionError as error:
        print(f"Local ASR setup failed ({error.code}): {error}", file=sys.stderr)
    except ImportError:
        print(
            "Required packages are missing; run with --install-dependencies first.", file=sys.stderr
        )
    except (OSError, ValueError, RuntimeError) as error:
        print(
            f"Local ASR setup failed: {type(error).__name__}. Check model files, WAV format, and connectivity for explicit downloads.",
            file=sys.stderr,
        )
    return 1


if __name__ == "__main__":
    raise SystemExit(main())