audio_transcriber.py

"""
audio_transcriber.py — Background Audio → SRT Transcription
Model: faster-whisper large-v3 (CPU, int8 quantization)

Fixes applied:
- word_timestamps=True for accurate per-word timing
- VAD threshold lowered to 0.20 to avoid missing quiet speech
- no_speech_threshold=0.5 (less aggressive segment dropping)
- Hallucination filter: skip segments with bad logprob or high compression
- Long segment split at punctuation boundaries (max 80 chars)
"""

import os
import re
import logging
import tempfile
import threading

from job_store import update_job, save_result

logger = logging.getLogger(__name__)

MODEL     = "large-v3"
MODEL_DIR = os.path.join(tempfile.gettempdir(), "whisper_models")

SUPPORTED_LANGUAGES = {
    "auto": "Auto-detect",
    "en":   "English",
    "ja":   "Japanese (日本語)",
    "ko":   "Korean (한국어)",
    "zh":   "Chinese (中文)",
    "th":   "Thai (ภาษาไทย)",
    "hi":   "Hindi (हिन्दी)",
    "fr":   "French",
    "de":   "German",
    "es":   "Spanish",
    "it":   "Italian",
    "pt":   "Portuguese",
    "ru":   "Russian",
    "ar":   "Arabic",
    "vi":   "Vietnamese",
    "id":   "Indonesian",
}

# ── Hallucination patterns Whisper commonly generates ─────────────────────
_HALLUCINATION_RE = re.compile(
    r"^\s*(\[.*?\]|\(.*?\)|♪.*?♪|thank you[.!]?|thanks for watching[.!]?)\s*$",
    re.IGNORECASE
)
_MAX_CHARS = 80   # max chars per subtitle line before splitting


# ─────────────────────────────────────────────────────────────────────────────
# Timestamp formatting
# ─────────────────────────────────────────────────────────────────────────────

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


# ─────────────────────────────────────────────────────────────────────────────
# Segment filtering — remove hallucinations
# ─────────────────────────────────────────────────────────────────────────────

def _is_hallucination(seg) -> bool:
    """
    Return True if this segment should be dropped.
    Whisper commonly hallucinates:
    - [Music] [Applause] (brackets)
    - Very low avg_logprob (AI is uncertain)
    - Very high compression_ratio (repetitive text)
    """
    text = seg.text.strip()

    # Empty
    if not text:
        return True

    # Known hallucination patterns
    if _HALLUCINATION_RE.match(text):
        return True

    # Low confidence — model very uncertain about this segment
    if hasattr(seg, 'avg_logprob') and seg.avg_logprob < -1.0:
        return True

    # High repetition — model is stuck in a loop
    if hasattr(seg, 'compression_ratio') and seg.compression_ratio > 2.4:
        return True

    # Timestamp sanity — end must be after start
    if seg.end <= seg.start:
        return True

    return False


# ─────────────────────────────────────────────────────────────────────────────
# Long segment splitting
# ─────────────────────────────────────────────────────────────────────────────

def _split_segment(seg) -> list:
    """
    If a segment is too long (>_MAX_CHARS), split it at sentence/phrase
    boundaries using word-level timestamps.

    Returns list of (start, end, text) tuples.
    If word timestamps are unavailable, returns the segment as-is.
    """
    text = seg.text.strip()

    # Short enough — no split needed
    if len(text) <= _MAX_CHARS:
        return [(seg.start, seg.end, text)]

    # Try word-level split
    words = getattr(seg, 'words', None)
    if not words:
        return [(seg.start, seg.end, text)]

    # Build sub-segments from word timestamps
    result     = []
    cur_words  = []
    cur_start  = words[0].start if words else seg.start

    for word in words:
        cur_words.append(word)
        cur_text = "".join(w.word for w in cur_words).strip()

        # Split at punctuation or when too long
        at_boundary = cur_text and cur_text[-1] in ".!?,;:…—"
        too_long    = len(cur_text) > _MAX_CHARS

        if (at_boundary or too_long) and cur_words:
            result.append((cur_start, word.end, cur_text))
            cur_words = []
            cur_start = word.end

    # Remainder
    if cur_words:
        cur_text = "".join(w.word for w in cur_words).strip()
        if cur_text:
            result.append((cur_start, seg.end, cur_text))

    return result if result else [(seg.start, seg.end, text)]


# ─────────────────────────────────────────────────────────────────────────────
# Segments → SRT
# ─────────────────────────────────────────────────────────────────────────────

def _segments_to_srt(segments: list) -> str:
    """
    Convert filtered segments to SRT string.
    - Drops hallucinations
    - Splits long segments at word boundaries
    - Ensures timestamp ordering (no backward jumps)
    """
    lines   = []
    idx     = 1
    prev_end = 0.0

    for seg in segments:
        # Drop bad segments
        if _is_hallucination(seg):
            logger.debug("Dropped hallucination: %s", seg.text[:50])
            continue

        # Split if too long
        sub_segs = _split_segment(seg)

        for (start, end, text) in sub_segs:
            text = text.strip()
            if not text:
                continue

            # Timestamp sanity: don't go backwards
            start = max(start, prev_end)
            if end <= start:
                end = start + 0.5   # give at least 0.5s

            lines += [str(idx), f"{_fmt(start)} --> {_fmt(end)}", text, ""]
            idx += 1
            prev_end = end

    return "\n".join(lines)


# ─────────────────────────────────────────────────────────────────────────────
# Background transcription worker
# ─────────────────────────────────────────────────────────────────────────────

def _transcribe_worker(job_id: str, audio_bytes: bytes,
                       filename: str, language: str) -> None:
    tmp_path = None
    try:
        from faster_whisper import WhisperModel

        def cb(frac, msg):
            update_job(job_id, progress=frac, status_msg=msg)

        cb(0.0, "Model load လုပ်နေသည် (large-v3)…")

        suffix = os.path.splitext(filename)[1] or ".mp3"
        fd, tmp_path = tempfile.mkstemp(suffix=suffix)
        os.write(fd, audio_bytes)
        os.close(fd)

        model = WhisperModel(
            MODEL, device="cpu", compute_type="int8",
            download_root=MODEL_DIR,
        )
        cb(0.05, "large-v3 load ပြီး — Transcribing…")

        lang_code = None if language == "auto" else language

        segments_gen, info = model.transcribe(
            tmp_path,
            language=lang_code,
            beam_size=5,
            temperature=0.0,
            condition_on_previous_text=False,   # reduce hallucination loops

            # ── Timestamp accuracy ─────────────────────────────────────────
            word_timestamps=True,               # per-word timing for accurate SRT

            # ── VAD — catch more speech, miss less ────────────────────────
            vad_filter=True,
            vad_parameters={
                "min_silence_duration_ms": 300,   # was 500 → catch short pauses
                "threshold": 0.20,               # was 0.35 → less aggressive
            },

            # ── Segment quality control ────────────────────────────────────
            no_speech_threshold=0.5,             # was default 0.6 → keep more
            log_prob_threshold=-1.0,             # drop very uncertain segments
            compression_ratio_threshold=2.4,     # drop repetitive segments
        )

        detected = info.language
        duration = max(info.duration or 1.0, 1.0)
        cb(0.08, f"ဘာသာစကား: {detected} · {duration:.0f}s — Segments ထုတ်နေသည်…")

        segments = []
        for seg in segments_gen:
            segments.append(seg)
            pct = min(0.08 + (seg.end / duration) * 0.90, 0.98)
            cb(pct, f"Transcribing… {seg.end:.0f}s/{duration:.0f}s  [{len(segments)} segs]")

        if not segments:
            update_job(job_id, status="failed",
                       error="Audio မှ speech မတွေ့ပါ",
                       status_msg="⚠ Speech မတွေ့ပါ")
            return

        cb(0.99, f"SRT ဖန်တီးနေသည်… ({len(segments)} segments)")
        srt_text  = _segments_to_srt(segments)
        srt_bytes = srt_text.encode("utf-8")

        # Count actual subtitle entries
        entry_count = srt_text.count("\n\n") + (1 if srt_text.strip() else 0)
        update_job(job_id, status_msg=f"ပြီးစီးသည် ✓  {entry_count} subtitles · {detected}")
        save_result(job_id, srt_bytes)

    except ImportError:
        update_job(job_id, status="failed",
                   error="faster-whisper မထည့်ရသေးပါ — requirements.txt စစ်ပါ",
                   status_msg="❌ Import Error")

    except Exception as exc:
        logger.exception("Transcription worker error")
        update_job(job_id, status="failed",
                   error=str(exc)[:400],
                   status_msg=f"❌ Error: {exc}")
    finally:
        if tmp_path and os.path.exists(tmp_path):
            try: os.unlink(tmp_path)
            except Exception: pass


# ─────────────────────────────────────────────────────────────────────────────
# Public API
# ─────────────────────────────────────────────────────────────────────────────

def run_transcription_bg(job_id: str, audio_bytes: bytes,
                         filename: str, language: str = "auto") -> threading.Thread:
    """Spawn background transcription thread. Survives browser disconnect."""
    update_job(job_id, status="running", progress=0.0,
               status_msg="Background thread started…")
    t = threading.Thread(
        target=_transcribe_worker,
        args=(job_id, audio_bytes, filename, language),
        daemon=True,
        name=f"whisper-{job_id}",
    )
    t.start()
    return t