app.py

"""
Gradio Space: Assyrian Neo-Aramaic (Urmi) ASR — upload or record, chunked inference.
"""
from __future__ import annotations

import os
import time
import warnings
from pathlib import Path
from typing import Callable, Literal

import gradio as gr
import librosa
import numpy as np
import torch
from transformers import Wav2Vec2BertForCTC, Wav2Vec2BertProcessor

from converters import russian_researchers_to_final_variant, russian_researchers_to_khan2016

warnings.filterwarnings("ignore", category=UserWarning)

MODEL_ID = os.environ.get("MODEL_ID", "Selest/wav2vec2-bert_Assyrian_Urmi_ASR_model")
TARGET_SR = 16_000

# Жёсткий потолок длины одного прохода (сек.) — длиннее режем по тишине внутри.
MAX_CHUNK_SEC = float(os.environ.get("MAX_CHUNK_SEC", "28"))
# Если пауза между двумя соседними VAD-сегментами не больше этого (мс), один сплошной
# вырезок wav[start:end] по таймлайну — короткая тишина между репликами остаётся в чанке.
# Поставьте 0, чтобы не объединять интервалы по таймлайну (склейка только на этапе pack до 28 с).
MERGE_GAP_MS = int(os.environ.get("MERGE_GAP_MS", "400"))
# При упаковке нескольких интервалов речи в один чанк до MAX_CHUNK_SEC — тишина между ними (мс).
PACK_JOIN_SILENCE_MS = int(os.environ.get("PACK_JOIN_SILENCE_MS", "400"))
# Минимальная длительность тишины при разрезе слишком длинного куска (мс).
MIN_SILENCE_CUT_MS = int(os.environ.get("MIN_SILENCE_CUT_MS", "150"))

# Silero VAD (torch.hub). VAD держим на CPU, чтобы не забирать VRAM у ASR.
SILERO_THRESHOLD = float(os.environ.get("SILERO_THRESHOLD", "0.5"))
SILERO_MIN_SPEECH_MS = int(os.environ.get("SILERO_MIN_SPEECH_MS", "250"))
SILERO_MIN_SILENCE_MS = int(os.environ.get("SILERO_MIN_SILENCE_MS", "100"))
SILERO_SPEECH_PAD_MS = int(os.environ.get("SILERO_SPEECH_PAD_MS", "64"))
# Второй проход Silero только по фрагменту, который всё ещё длиннее MAX_CHUNK_SEC.
SILERO_REFINE_ENABLED = os.environ.get("SILERO_REFINE_ENABLED", "1").lower() not in (
    "0",
    "false",
    "no",
)
SILERO_REFINE_MIN_SILENCE_MS = int(
    os.environ.get("SILERO_REFINE_MIN_SILENCE_MS", "80")
)

OutputSystem = Literal["researchers", "khan", "cyrillic"]

_DEVICE: torch.device | None = None
_PROCESSOR: Wav2Vec2BertProcessor | None = None
_MODEL: Wav2Vec2BertForCTC | None = None
_SILERO_MODEL: torch.nn.Module | None = None
_SILERO_GET_SPEECH_TIMESTAMPS: Callable[..., list[dict[str, int]]] | None = None


def _get_device() -> torch.device:
    global _DEVICE
    if _DEVICE is None:
        _DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    return _DEVICE


def load_model() -> None:
    global _PROCESSOR, _MODEL
    if _MODEL is not None:
        return
    device = _get_device()
    dtype = torch.float16 if device.type == "cuda" else torch.float32
    _PROCESSOR = Wav2Vec2BertProcessor.from_pretrained(MODEL_ID)
    _MODEL = Wav2Vec2BertForCTC.from_pretrained(MODEL_ID, torch_dtype=dtype)
    _MODEL.to(device)
    _MODEL.eval()


def load_silero_vad() -> None:
    """Silero VAD через torch.hub (нужен интернет при первом запуске / на Space)."""
    global _SILERO_MODEL, _SILERO_GET_SPEECH_TIMESTAMPS
    if _SILERO_MODEL is not None:
        return
    model, utils = torch.hub.load(
        "snakers4/silero-vad",
        "silero_vad",
        force_reload=False,
        onnx=False,
        trust_repo=True,
        skip_validation=True,
    )
    vad_dev = torch.device("cpu")
    model = model.to(vad_dev)
    model.eval()
    _SILERO_MODEL = model
    get_speech_timestamps = utils[0]
    _SILERO_GET_SPEECH_TIMESTAMPS = get_speech_timestamps


def _ensure_mono_16k(wav: np.ndarray, sr: int) -> np.ndarray:
    if wav.ndim > 1:
        wav = np.mean(wav, axis=1)
    wav = np.asarray(wav, dtype=np.float32)
    if sr != TARGET_SR:
        wav = librosa.resample(wav, orig_sr=sr, target_sr=TARGET_SR)
    # Нормализация в типичный диапазон [-1, 1]
    peak = float(np.max(np.abs(wav))) if wav.size else 0.0
    if peak > 1.0:
        wav = wav / peak
    return wav


def _vad_boundary_to_int(x: object) -> int:
    if isinstance(x, torch.Tensor):
        return int(x.item())
    return int(x)


def _merge_close_interval_pairs(
    pairs: list[tuple[int, int]], merge_gap_samples: int
) -> list[tuple[int, int]]:
    """Объединяет соседние по времени интервалы [start, end), если зазор между ними короткий."""
    if not pairs:
        return []
    ordered = sorted((int(s), int(e)) for s, e in pairs)
    out: list[tuple[int, int]] = [ordered[0]]
    for s, e in ordered[1:]:
        ps, pe = out[-1]
        if s - pe <= merge_gap_samples:
            out[-1] = (ps, max(pe, e))
        else:
            out.append((s, e))
    return out


def _silero_speech_intervals(wav: np.ndarray) -> list[tuple[int, int]]:
    """Интервалы речи по Silero VAD (сэмплы 16 kHz)."""
    load_silero_vad()
    if _SILERO_MODEL is None or _SILERO_GET_SPEECH_TIMESTAMPS is None:
        raise RuntimeError("Silero VAD is not loaded")

    n = int(wav.size)
    if n == 0:
        return []

    w = torch.from_numpy(np.ascontiguousarray(wav, dtype=np.float32))
    timestamps = _SILERO_GET_SPEECH_TIMESTAMPS(
        w,
        _SILERO_MODEL,
        threshold=SILERO_THRESHOLD,
        sampling_rate=TARGET_SR,
        min_speech_duration_ms=SILERO_MIN_SPEECH_MS,
        max_speech_duration_s=MAX_CHUNK_SEC,
        min_silence_duration_ms=SILERO_MIN_SILENCE_MS,
        speech_pad_ms=SILERO_SPEECH_PAD_MS,
        return_seconds=False,
    )
    raw_pairs: list[tuple[int, int]] = []
    for t in timestamps:
        s = _vad_boundary_to_int(t["start"])
        e = _vad_boundary_to_int(t["end"])
        s = int(max(0, min(n, s)))
        e = int(max(0, min(n, e)))
        if e > s:
            raw_pairs.append((s, e))

    ordered = sorted(raw_pairs, key=lambda p: (p[0], p[1]))
    if MERGE_GAP_MS <= 0:
        return ordered
    merge_gap = int(TARGET_SR * (MERGE_GAP_MS / 1000.0))
    return _merge_close_interval_pairs(ordered, merge_gap)


def _silero_refine_threshold() -> float:
    raw = os.environ.get("SILERO_REFINE_THRESHOLD")
    if raw is None or str(raw).strip() == "":
        return SILERO_THRESHOLD
    return float(raw)


def _silero_refine_oversized_chunk(y: np.ndarray) -> list[np.ndarray] | None:
    """
    Повторный Silero на одном длинном вырезке. Интервалы не склеиваем, чтобы не ломать разрез.
    """
    if not SILERO_REFINE_ENABLED:
        return None
    load_silero_vad()
    if _SILERO_MODEL is None or _SILERO_GET_SPEECH_TIMESTAMPS is None:
        return None

    n = int(y.size)
    if n == 0:
        return None

    w = torch.from_numpy(np.ascontiguousarray(y, dtype=np.float32))
    timestamps = _SILERO_GET_SPEECH_TIMESTAMPS(
        w,
        _SILERO_MODEL,
        threshold=_silero_refine_threshold(),
        sampling_rate=TARGET_SR,
        min_speech_duration_ms=SILERO_MIN_SPEECH_MS,
        max_speech_duration_s=MAX_CHUNK_SEC,
        min_silence_duration_ms=SILERO_REFINE_MIN_SILENCE_MS,
        speech_pad_ms=SILERO_SPEECH_PAD_MS,
        return_seconds=False,
    )
    pieces: list[np.ndarray] = []
    for t in timestamps:
        s = _vad_boundary_to_int(t["start"])
        e = _vad_boundary_to_int(t["end"])
        s = int(max(0, min(n, s)))
        e = int(max(0, min(n, e)))
        if e > s:
            pieces.append(np.ascontiguousarray(y[s:e], dtype=np.float32))

    if not pieces:
        return None
    if len(pieces) == 1 and pieces[0].size >= int(n * 0.99):
        return None
    return pieces


def _frame_rms_energy(y: np.ndarray, frame_len: int, hop: int) -> tuple[np.ndarray, np.ndarray]:
    """RMS по коротким кадрам; возвращает rms и индекс центра каждого кадра."""
    if y.size < frame_len:
        return np.array([], dtype=np.float32), np.array([], dtype=np.int64)
    n_frames = 1 + (y.size - frame_len) // hop
    rms = np.empty(n_frames, dtype=np.float32)
    centers = np.empty(n_frames, dtype=np.int64)
    for i in range(n_frames):
        start = i * hop
        frame = y[start : start + frame_len]
        rms[i] = float(np.sqrt(np.mean(frame * frame)))
        centers[i] = start + frame_len // 2
    return rms, centers


def _split_long_segment_at_silence(y: np.ndarray, sr: int) -> tuple[list[np.ndarray], bool]:
    """
    Делит длинный сегмент по самой длинной внутренней паузе.
    Если подходящей паузы нет — один кусок (вызывающий код при необходимости режет принудительно).
    Returns (parts, found_silence).
    """
    max_samples = int(sr * MAX_CHUNK_SEC)
    if y.size <= max_samples:
        return [y], True

    frame_len = int(0.02 * sr)
    hop = int(0.010 * sr)
    min_sil_frames = max(2, int(MIN_SILENCE_CUT_MS / 10))

    rms, centers = _frame_rms_energy(y, frame_len, hop)
    if rms.size == 0:
        mid = y.size // 2
        return [y[:mid], y[mid:]], False

    thresh = float(np.percentile(rms, 22))
    if thresh <= 1e-8:
        thresh = float(np.max(rms) * 0.15)
    silent = rms < thresh

    best_run = 0
    best_start = -1
    cur_start = -1
    cur_len = 0
    for i, s in enumerate(silent):
        if s:
            if cur_start < 0:
                cur_start = i
                cur_len = 1
            else:
                cur_len += 1
        else:
            if cur_len > 0 and cur_len >= min_sil_frames and cur_len > best_run:
                best_run = cur_len
                best_start = cur_start
            cur_start = -1
            cur_len = 0
    if cur_len > 0 and cur_len >= min_sil_frames and cur_len > best_run:
        best_run = cur_len
        best_start = cur_start

    if best_start < 0 or best_run < min_sil_frames:
        return [y], False

    cut_frame = best_start + best_run // 2
    cut_sample = int(centers[min(cut_frame, len(centers) - 1)])
    cut_sample = max(frame_len, min(y.size - frame_len, cut_sample))

    left, right = y[:cut_sample], y[cut_sample:]
    parts: list[np.ndarray] = []
    for half in (left, right):
        sub, _ = _split_long_segment_at_silence(half, sr)
        parts.extend(sub)
    return parts, True


def _force_split_midpoint(y: np.ndarray) -> list[np.ndarray]:
    """Последний резерв: середина по времени (есть риск порезать слово)."""
    mid = max(1, y.size // 2)
    return [y[:mid], y[mid:]]


def _shrink_to_max_sec(y: np.ndarray, sr: int) -> list[np.ndarray]:
    """
    Длина ≤ MAX_CHUNK_SEC: сначала второй проход Silero на саб-фрагменте,
    затем разрез по RMS, затем пополам.
    """
    max_samples = int(sr * MAX_CHUNK_SEC)
    if y.size <= max_samples or y.size == 0:
        return [y] if y.size > 0 else []

    refined = _silero_refine_oversized_chunk(y)
    if refined is not None:
        if len(refined) >= 2:
            acc: list[np.ndarray] = []
            for p in refined:
                acc.extend(_shrink_to_max_sec(p, sr))
            return acc
        if len(refined) == 1 and refined[0].size < y.size:
            return _shrink_to_max_sec(refined[0], sr)

    parts, _used = _split_long_segment_at_silence(y, sr)
    if len(parts) == 1 and parts[0].size == y.size:
        a, b = _force_split_midpoint(y)
        return _shrink_to_max_sec(a, sr) + _shrink_to_max_sec(b, sr)

    acc2: list[np.ndarray] = []
    for p in parts:
        acc2.extend(_shrink_to_max_sec(p, sr))
    return acc2


def _pack_speech_arrays_under_limit(
    speech_arrays: list[np.ndarray],
    max_samples: int,
    join_pad_samples: int,
) -> list[np.ndarray]:
    """
    Жадно склеивает подряд идущие интервалы речи: между каждыми двумя вставляет паузу join_pad_samples.

    Пока возможно, держим сумму ≤ max_samples; если без паузы влезает в лимит, а с паузой —
    уже нет, всё равно склеиваем с паузой в один буфер (дальше _shrink_to_max_sec разрежет при нужде).
    Если даже без паузы не влезает — закрываем текущий чанк и начинаем новый с y.
    """
    pad = max(0, int(join_pad_samples))
    gap = np.zeros(pad, dtype=np.float32) if pad > 0 else None

    out: list[np.ndarray] = []
    cur: np.ndarray | None = None
    for y in speech_arrays:
        if y.size == 0:
            continue
        y = np.ascontiguousarray(y, dtype=np.float32)
        if cur is None:
            cur = y
            continue
        extra = pad + int(y.size)
        ny = int(y.size)
        if cur.size + extra <= max_samples:
            if gap is not None:
                cur = np.concatenate([cur, gap, y])
            else:
                cur = np.concatenate([cur, y])
        elif pad > 0 and gap is not None and cur.size + ny <= max_samples:
            cur = np.concatenate([cur, gap, y])
        else:
            out.append(cur)
            cur = y
    if cur is not None:
        out.append(cur)
    return out


def build_speech_chunks(wav: np.ndarray, sr: int) -> list[np.ndarray]:
    """Silero → при MERGE_GAP_MS>0 короткие паузы между репликами входят в один интервал → pack до MAX_CHUNK_SEC → разрез."""
    if sr != TARGET_SR:
        wav = librosa.resample(wav, orig_sr=sr, target_sr=TARGET_SR)
        sr = TARGET_SR

    base_intervals = _silero_speech_intervals(wav)
    if not base_intervals:
        raw_chunks = [wav]
    else:
        raw_chunks = [
            np.ascontiguousarray(wav[s:e], dtype=np.float32)
            for s, e in base_intervals
            if e > s
        ]

    max_samples = int(sr * MAX_CHUNK_SEC)
    join_pad_samples = max(0, int(sr * PACK_JOIN_SILENCE_MS / 1000.0))
    packed = (
        _pack_speech_arrays_under_limit(raw_chunks, max_samples, join_pad_samples)
        if base_intervals
        else raw_chunks
    )

    final: list[np.ndarray] = []
    for chunk in packed:
        if chunk.size == 0:
            continue
        final.extend(_shrink_to_max_sec(chunk, sr))

    if not final:
        final = [wav]
    return final


def _apply_output_system(text: str, system: OutputSystem) -> str:
    t = (text or "").strip()
    if not t:
        return t
    if system == "researchers":
        return t
    if system == "khan":
        return russian_researchers_to_khan2016(t)
    return russian_researchers_to_final_variant(t)


def _transcribe_array(wav: np.ndarray) -> str:
    assert _PROCESSOR is not None and _MODEL is not None
    device = _get_device()
    inputs = _PROCESSOR(
        wav, sampling_rate=TARGET_SR, return_tensors="pt", padding=True
    )
    inputs = {k: v.to(device) for k, v in inputs.items()}
    with torch.inference_mode():
        logits = _MODEL(**inputs).logits
    pred_ids = torch.argmax(logits, dim=-1)
    return (_PROCESSOR.batch_decode(pred_ids)[0] or "").strip()


def _format_eta(seconds: float) -> str:
    if seconds < 60:
        return f"{max(0, seconds):.0f} s"
    m, s = divmod(int(seconds), 60)
    return f"{m} min {s} s"


def _fmt_audio_sec(x: float) -> str:
    return f"{max(0.0, float(x)):.1f}"


def _input_audio_to_path(audio_in: object) -> str | None:
    """Resolve path from Gradio 6: str, Path, FileData-like, or dict with path."""
    if audio_in is None:
        return None
    if isinstance(audio_in, str):
        s = audio_in.strip()
        return s if s else None
    if isinstance(audio_in, Path):
        s = str(audio_in)
        return s if s else None
    if isinstance(audio_in, dict):
        p = audio_in.get("path") or audio_in.get("name")
        return str(p) if p else None
    p = getattr(audio_in, "path", None)
    if p is not None:
        return str(p)
    return None


def transcribe(
    audio_in: object,
    output_system: str,
    progress: gr.Progress | None = None,
):
    load_model()
    load_silero_vad()
    audio_path = _input_audio_to_path(audio_in)
    if not audio_path or not os.path.isfile(audio_path):
        yield "", "Upload or record audio (16 kHz preferred)."
        return

    wav, sr = librosa.load(audio_path, sr=None, mono=True)
    wav = _ensure_mono_16k(wav, int(sr))

    chunks = build_speech_chunks(wav, TARGET_SR)
    n = len(chunks)
    chunk_lens_sec = [float(c.size) / TARGET_SR for c in chunks]
    total_work_sec = float(sum(chunk_lens_sec))
    system_key: OutputSystem = output_system  # type: ignore[assignment]

    texts: list[str] = []
    chunk_times: list[float] = []
    t_wall0 = time.perf_counter()

    def _progress_by_audio(processed_audio_sec: float) -> float:
        if total_work_sec <= 0:
            return 1.0 if n == 0 else 0.0
        return min(1.0, max(0.0, processed_audio_sec / total_work_sec))

    for i, chunk in enumerate(chunks):
        done = i + 1
        cur_sec = chunk_lens_sec[i] if i < len(chunk_lens_sec) else 0.0
        processed_before_sec = sum(chunk_lens_sec[:i])

        eta_str = "…"
        if chunk_times:
            avg = sum(chunk_times) / len(chunk_times)
            eta_str = _format_eta(avg * (n - done))

        processed_mid = processed_before_sec + 0.35 * cur_sec
        if progress is not None:
            progress(
                _progress_by_audio(processed_mid),
                desc=(
                    f"Audio ~{_fmt_audio_sec(processed_mid)}/{_fmt_audio_sec(total_work_sec)} s · "
                    f"segment {done}/{n} · ~{eta_str} left"
                ),
            )

        t0 = time.perf_counter()
        part = _transcribe_array(chunk)
        chunk_times.append(time.perf_counter() - t0)
        if part:
            texts.append(part)

        processed_after_sec = sum(chunk_lens_sec[:done])

        if processed_after_sec > 1e-6:
            elapsed = time.perf_counter() - t_wall0
            remaining_audio = max(0.0, total_work_sec - processed_after_sec)
            eta_after = _format_eta((elapsed / processed_after_sec) * remaining_audio)
        elif chunk_times:
            avg = sum(chunk_times) / len(chunk_times)
            remaining_sec = avg * (n - done)
            eta_after = _format_eta(remaining_sec)
        else:
            eta_after = "…"

        if done >= n:
            eta_after = "complete"

        prog_tail = (
            f"~{eta_after} remaining"
            if done < n
            else "complete"
        )
        if progress is not None:
            progress(
                _progress_by_audio(processed_after_sec),
                desc=(
                    f"Audio ~{_fmt_audio_sec(processed_after_sec)}/{_fmt_audio_sec(total_work_sec)} s · "
                    f"segment {done}/{n} · {prog_tail}"
                ),
            )

        raw_merged = " ".join(texts).strip()
        if eta_after == "complete":
            status_time_line = "Complete."
        else:
            status_time_line = f"ETA: ~{eta_after} remaining."
        status_mid = (
            f"Processed ~{_fmt_audio_sec(processed_after_sec)} s of ~{_fmt_audio_sec(total_work_sec)} s "
            f"(sum of speech segments; long file silences excluded). "
            f"Segment {done}/{n}. {status_time_line} "
            f"Intermediate text: raw model output (Russian researchers convention)."
        )
        yield raw_merged, status_mid

        if _get_device().type == "cuda":
            torch.cuda.empty_cache()

    merged = " ".join(texts).strip()
    merged_out = _apply_output_system(merged, system_key)

    wall_sec = time.perf_counter() - t_wall0
    status = (
        f"Done: {n} segment(s), ~{_fmt_audio_sec(total_work_sec)} s of audio transcribed, "
        f"wall time ~{wall_sec:.1f} s."
    )
    yield merged_out, status


def build_demo() -> gr.Blocks:
    # Model and Silero load on first submit so Space startup does not hit the health-check timeout.
    with gr.Blocks(title="Assyrian Urmi ASR") as blocks:
        with gr.Row():
            audio_in = gr.Audio(
                sources=["upload", "microphone"],
                type="filepath",
                format="wav",
                label="Upload or record",
            )
        out_system = gr.Radio(
            choices=[
                ("Russian researchers (as returned by the model)", "researchers"),
                ("Khan (2016)", "khan"),
                ("Cyrillic letters", "cyrillic"),
            ],
            value="researchers",
            label="Output notation",
        )
        btn = gr.Button("Transcribe")
        out_text = gr.Textbox(label="Transcript", lines=8)
        status = gr.Textbox(label="Status / progress", lines=2)

        btn.click(
            fn=transcribe,
            inputs=[audio_in, out_system],
            outputs=[out_text, status],
        )

    return blocks.queue()


demo = build_demo()


if __name__ == "__main__":
    demo.launch()