second-brain/src/second_brain/transcribe.py

"""Audio/video transcription via faster-whisper (CTranslate2 backend).

Lazy-imports `faster_whisper` so the dev side — which doesn't transcribe
anything — doesn't need the heavyweight CUDA wheels installed.

Config keys (read from `[whisper]` in settings.toml, env-overrides win):
- `whisper_model`         : the model name. Default "large-v3".
- `whisper_device`        : "cuda" | "cpu" | "auto". Default "cuda".
- `whisper_compute_type`  : "float16" | "int8" | "int8_float16" | …
                            Default float16 on cuda, int8 on cpu.

Outputs an SRT file next to `config.subtitles_dir/{stem}.srt` plus a
plain-text transcript string. The pipeline writes the text into
`sources.transcript_text` and the path into `sources.transcript_path`,
matching the previous torch-whisper contract so downstream extract/embed
doesn't need to change.
"""

from __future__ import annotations

import logging
import os
from pathlib import Path
from typing import Iterable, Optional

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Config resolution
# ---------------------------------------------------------------------------


def _resolve(cfg: dict | None, env_key: str, cfg_key: str, default: str) -> str:
    env = os.environ.get(env_key)
    if env:
        return env
    if cfg is not None:
        v = cfg.get(cfg_key)
        if v:
            return v
    return default


def resolve_settings(config_block: dict | None = None) -> dict[str, str]:
    """Resolve whisper runtime settings from env > config block > defaults.

    Pulls model/device/compute_type and picks a sane compute_type per device
    when none was specified. Returns plain strings so the caller can stash
    them in logs / pass straight to faster-whisper.
    """
    model = _resolve(config_block, "WHISPER_MODEL", "whisper_model", "large-v3")
    device = _resolve(config_block, "WHISPER_DEVICE", "whisper_device", "cuda")

    explicit_compute = None
    if config_block is not None:
        explicit_compute = config_block.get("whisper_compute_type")
    explicit_compute = os.environ.get("WHISPER_COMPUTE_TYPE", explicit_compute)

    if explicit_compute:
        compute_type = explicit_compute
    elif device == "cpu":
        compute_type = "int8"
    else:
        compute_type = "float16"

    return {"model": model, "device": device, "compute_type": compute_type}


# ---------------------------------------------------------------------------
# Transcriber
# ---------------------------------------------------------------------------


class Transcriber:
    """faster-whisper wrapper. Model is loaded lazily on first transcribe."""

    def __init__(
        self,
        model: str = "large-v3",
        device: str = "cuda",
        compute_type: str = "float16",
    ) -> None:
        self.model = model
        self.device = device
        self.compute_type = compute_type
        self._impl = None

    def _load(self) -> None:
        if self._impl is not None:
            return
        try:
            from faster_whisper import WhisperModel  # type: ignore[import-not-found]
        except ImportError as exc:
            raise RuntimeError(
                "faster-whisper is not installed. On the tower run: "
                "uv sync --extra tower"
            ) from exc

        logger.info(
            "loading faster-whisper model=%s device=%s compute_type=%s",
            self.model,
            self.device,
            self.compute_type,
        )
        self._impl = WhisperModel(
            self.model, device=self.device, compute_type=self.compute_type
        )

    def transcribe(
        self,
        media_file: Path,
        *,
        language: Optional[str] = None,
        beam_size: int = 5,
    ) -> tuple[str, list[dict]]:
        """Transcribe `media_file`. Returns (full_text, segments).

        Each segment dict has start/end (float seconds) and text (str) —
        same shape the SRT writer needs.
        """
        self._load()
        assert self._impl is not None  # for the type-checker

        segments_iter, info = self._impl.transcribe(
            str(media_file),
            language=language,
            beam_size=beam_size,
        )
        logger.info(
            "transcribe start: %s language=%s duration=%.1fs",
            media_file.name,
            getattr(info, "language", "?"),
            getattr(info, "duration", 0.0) or 0.0,
        )

        segments: list[dict] = []
        text_parts: list[str] = []
        for seg in segments_iter:
            t = (seg.text or "").strip()
            segments.append({"start": float(seg.start), "end": float(seg.end), "text": t})
            if t:
                text_parts.append(t)

        return " ".join(text_parts), segments


# ---------------------------------------------------------------------------
# SRT writer
# ---------------------------------------------------------------------------


def write_srt(segments: Iterable[dict], output: Path) -> None:
    """Write a list of {start,end,text} dicts as SRT into `output`."""
    output.parent.mkdir(parents=True, exist_ok=True)
    with open(output, "w", encoding="utf-8") as fh:
        for i, seg in enumerate(segments, start=1):
            start = _fmt_timestamp(seg["start"])
            end = _fmt_timestamp(seg["end"])
            fh.write(f"{i}\n{start} --> {end}\n{seg['text']}\n\n")


def _fmt_timestamp(seconds: float) -> str:
    h = int(seconds // 3600)
    m = int((seconds % 3600) // 60)
    s = int(seconds % 60)
    ms = int((seconds % 1) * 1000)
    return f"{h:02d}:{m:02d}:{s:02d},{ms:03d}"


__all__ = ["Transcriber", "resolve_settings", "write_srt"]