diarize.py

"""Bước 2 — Phân tích người nói (diarization).

Wrapper gọn quanh pyannote: audio → danh sách `SpeakerTurn` → gán `cue.speaker`
theo độ phủ thời gian lớn nhất. pyannote chỉ import khi thực sự chạy (lazy) để
Bước 1 vẫn dùng được local không cần GPU.
"""

from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
from typing import List, Optional

from .srt import Cue


@dataclass
class SpeakerTurn:
    """Một lượt nói liên tục của một speaker."""

    speaker: str
    start: float
    end: float

    @property
    def duration(self) -> float:
        return max(0.0, self.end - self.start)


def diarize_audio(
    audio_path: str | Path,
    *,
    model_name: str = "pyannote/speaker-diarization-community-1",
    hf_token: Optional[str] = None,
    num_speakers: Optional[int] = None,
    max_time: Optional[float] = None,
) -> List[SpeakerTurn]:
    """Chạy pyannote trên audio (16 kHz mono khuyến nghị) → SpeakerTurn.

    `max_time`: chỉ diarize tới mốc này (vd max(cue.end)) để tiết kiệm thời gian.
    """
    try:
        import torch
        from pyannote.audio import Pipeline
    except ImportError as e:  # pragma: no cover
        raise RuntimeError(
            "Bước 2 cần `pyannote.audio` + `torch`. Cài qua requirements-colab.txt."
        ) from e

    pipeline = Pipeline.from_pretrained(model_name, token=hf_token)
    if torch.cuda.is_available():
        pipeline.to(torch.device("cuda"))

    params = {}
    if num_speakers is not None:
        params["num_speakers"] = num_speakers

    diarization = pipeline(str(audio_path), **params)

    turns: List[SpeakerTurn] = []
    for segment, _, speaker in diarization.itertracks(yield_label=True):
        if max_time is not None and segment.start > max_time:
            continue
        turns.append(
            SpeakerTurn(speaker=str(speaker), start=float(segment.start), end=float(segment.end))
        )
    turns.sort(key=lambda t: t.start)
    return turns


def assign_speakers(cues: List[Cue], turns: List[SpeakerTurn]) -> List[Cue]:
    """Gán `cue.speaker` = speaker phủ thời gian cue nhiều nhất."""
    if not turns:
        return cues
    for cue in cues:
        best_speaker: Optional[str] = None
        best_overlap = 0.0
        for t in turns:
            overlap = min(cue.end, t.end) - max(cue.start, t.start)
            if overlap > best_overlap:
                best_overlap = overlap
                best_speaker = t.speaker
        if best_speaker is not None:
            cue.speaker = best_speaker
    return cues