scripts/progress.py - master - Max_Cherep/super-transcription

scripts/progress.py
249 строк8 KB
maxim.cherepanov
Show per-stage processing times in the jobs table.
01 авг 2026, 13:48
01 авг 2026, 13:48f44c127
"""Прогресс обработки: консоль + SQLite (throttle)."""

from __future__ import annotations

import sys
import time
from typing import Any

from db import update_job
from job_timings import dumps_stage_timings

BAR_LEN = 40
DB_THROTTLE_SEC = 1.0
TIMING_FLUSH_SEC = 2.0

WEIGHT_EXTRACT = 0.10
WEIGHT_ASR = 0.80
WEIGHT_DIARIZE = 0.05
WEIGHT_REPAIR = 0.05
WEIGHT_OCR = 0.05
WEIGHT_EXPORT = 0.10
WEIGHT_ASR_WITH_DIAR = 0.75
WEIGHT_ASR_WITH_DIAR_REPAIR = 0.70
WEIGHT_ASR_WITH_OCR = 0.75
WEIGHT_ASR_WITH_DIAR_OCR = 0.70
WEIGHT_ASR_WITH_DIAR_REPAIR_OCR = 0.65

_STAGE_SHORT = {
    "Извлечение аудио": "аудио",
    "Транскрибация": "ASR",
    "Диаризация": "diar",
    "Правка спикеров": "repair",
    "Скриншоты": "shots",
    "Экспорт": "экспорт",
}

_db_state: dict[int, dict[str, Any]] = {}


def clear_job_progress_state(job_id: int) -> None:
    _db_state.pop(job_id, None)


def update_job_progress(
    job_id: int,
    pct: float,
    detail: str,
    *,
    force: bool = False,
) -> None:
    pct = max(0.0, min(100.0, pct))
    state = _db_state.setdefault(job_id, {"t": 0.0, "detail": "", "pct": -1.0})
    now = time.monotonic()
    changed_detail = detail != state["detail"]
    if not force:
        if now - state["t"] < DB_THROTTLE_SEC and not changed_detail:
            if abs(pct - state["pct"]) < 1.0 and pct < 99.9:
                return
    update_job(job_id, progress_pct=pct, progress_detail=detail)
    state["t"] = now
    state["detail"] = detail
    state["pct"] = pct


class ProgressReporter:
    """Консольный бар + опционально запись в БД."""

    def __init__(self, job_id: int | None = None, *, prefix: str = "") -> None:
        self.job_id = job_id
        self.prefix = prefix
        self._last_line = ""
        self._tty = sys.stderr.isatty()

    def report(self, pct: float, detail: str, *, force_db: bool = False) -> None:
        pct = max(0.0, min(100.0, pct))
        self._console(pct, detail)
        if self.job_id is not None:
            update_job_progress(self.job_id, pct, detail, force=force_db)

    def _console(self, pct: float, detail: str) -> None:
        filled = int(BAR_LEN * pct / 100)
        bar = "█" * filled + "░" * (BAR_LEN - filled)
        head = f"{self.prefix}| " if self.prefix else "  "
        line = f"{head}[{bar}] {pct:5.1f}% | {detail}"
        if self._tty:
            print(f"\r{line}", end="", file=sys.stderr, flush=True)
            self._last_line = line
            return
        if line != self._last_line:
            print(line, file=sys.stderr, flush=True)
            self._last_line = line

    def finish_line(self) -> None:
        if self._tty and self._last_line:
            print(file=sys.stderr, flush=True)
            self._last_line = ""


def _normalize_weights(stages: list[tuple[str, float]]) -> list[tuple[str, float]]:
    total = sum(w for _, w in stages)
    if total <= 0:
        return stages
    return [(name, w / total) for name, w in stages]


class JobProgress:
    """Доли этапов пайплайна → общий %."""

    def __init__(
        self,
        reporter: ProgressReporter | None,
        *,
        include_diarize: bool = False,
        include_repair: bool = False,
        include_ocr: bool = False,
    ) -> None:
        self.reporter = reporter
        if include_diarize and include_repair and include_ocr:
            stages = [
                ("Извлечение аудио", WEIGHT_EXTRACT),
                ("Транскрибация", WEIGHT_ASR_WITH_DIAR_REPAIR_OCR),
                ("Диаризация", WEIGHT_DIARIZE),
                ("Правка спикеров", WEIGHT_REPAIR),
                ("Скриншоты", WEIGHT_OCR),
                ("Экспорт", WEIGHT_EXPORT),
            ]
        elif include_diarize and include_repair:
            stages = [
                ("Извлечение аудио", WEIGHT_EXTRACT),
                ("Транскрибация", WEIGHT_ASR_WITH_DIAR_REPAIR),
                ("Диаризация", WEIGHT_DIARIZE),
                ("Правка спикеров", WEIGHT_REPAIR),
                ("Экспорт", WEIGHT_EXPORT),
            ]
        elif include_diarize and include_ocr:
            stages = [
                ("Извлечение аудио", WEIGHT_EXTRACT),
                ("Транскрибация", WEIGHT_ASR_WITH_DIAR_OCR),
                ("Диаризация", WEIGHT_DIARIZE),
                ("Скриншоты", WEIGHT_OCR),
                ("Экспорт", WEIGHT_EXPORT),
            ]
        elif include_diarize:
            stages = [
                ("Извлечение аудио", WEIGHT_EXTRACT),
                ("Транскрибация", WEIGHT_ASR_WITH_DIAR),
                ("Диаризация", WEIGHT_DIARIZE),
                ("Экспорт", WEIGHT_EXPORT),
            ]
        elif include_ocr:
            stages = [
                ("Извлечение аудио", WEIGHT_EXTRACT),
                ("Транскрибация", WEIGHT_ASR_WITH_OCR),
                ("Скриншоты", WEIGHT_OCR),
                ("Экспорт", WEIGHT_EXPORT),
            ]
        else:
            stages = [
                ("Извлечение аудио", WEIGHT_EXTRACT),
                ("Транскрибация", WEIGHT_ASR),
                ("Экспорт", WEIGHT_EXPORT),
            ]
        self._stages = _normalize_weights(stages)
        self._cursor = 0
        self._base = 0.0
        self.timings: list[dict[str, Any]] = []
        self._step_t0 = time.monotonic()
        self._timing_flush_t = 0.0

    def _stage_short(self, name: str) -> str:
        return _STAGE_SHORT.get(name, name)

    def _persist_timings(self, *, live: dict[str, Any] | None = None) -> None:
        if self.reporter is None or self.reporter.job_id is None:
            return
        items = list(self.timings)
        if live is not None:
            items.append(live)
        update_job(
            self.reporter.job_id,
            stage_timings=dumps_stage_timings(items),
        )

    def set_step_progress(self, frac: float, label: str) -> None:
        frac = max(0.0, min(1.0, frac))
        if self._cursor >= len(self._stages):
            return
        name, weight = self._stages[self._cursor]
        detail = label or name
        pct = (self._base + weight * frac) * 100.0
        if self.reporter is not None:
            self.reporter.report(pct, detail)
        now = time.monotonic()
        if now - self._timing_flush_t >= TIMING_FLUSH_SEC:
            self._timing_flush_t = now
            elapsed = max(0.0, now - self._step_t0)
            self._persist_timings(
                live={
                    "name": self._stage_short(name),
                    "sec": round(elapsed, 1),
                    "live": True,
                }
            )

    def complete_step(self, label: str | None = None) -> None:
        if self._cursor >= len(self._stages):
            return
        name, weight = self._stages[self._cursor]
        elapsed = max(0.0, time.monotonic() - self._step_t0)
        self.timings.append(
            {"name": self._stage_short(name), "sec": round(elapsed, 1)}
        )
        self._base += weight
        self._cursor += 1
        self._step_t0 = time.monotonic()
        self._timing_flush_t = 0.0
        detail = label or name
        if self.reporter is not None:
            self.reporter.report(self._base * 100.0, detail, force_db=True)
            self._persist_timings()


def build_job_progress(
    reporter: ProgressReporter | None,
    opts: dict[str, Any] | None = None,
) -> JobProgress:
    speech_mode = "monologue"
    if opts:
        speech_mode = str(opts.get("speech_mode") or "monologue")
    from diarization import wants_diarization
    from dialogue_repair import repair_enabled

    include_diarize = wants_diarization(speech_mode)
    # шаг прогресса: UI dialogue_repair + dialogue + kill-switch
    want_repair = bool(opts.get("dialogue_repair")) if opts else False
    include_repair = (
        include_diarize
        and speech_mode == "dialogue"
        and want_repair
        and repair_enabled()
    )
    shots = (opts or {}).get("screenshots") or []
    include_ocr = isinstance(shots, list) and len(shots) > 0
    return JobProgress(
        reporter,
        include_diarize=include_diarize,
        include_repair=include_repair,
        include_ocr=include_ocr,
    )