models/transcriber.py

"""
ShortSmith v4 - Transcriber

Whisper-based audio transcription with chorus/repeat detection.
"""

from pathlib import Path
from typing import List, Dict, Optional
import difflib

from utils.logger import get_logger
import config as cfg

logger = get_logger("models.transcriber")


class Transcriber:
    """
    Transcribes audio using OpenAI Whisper (local, no API key needed).
    Also detects repeated lyric sections (chorus) from the transcript.
    """

    def __init__(self, model_name: Optional[str] = None):
        self.model_name = model_name or cfg.WHISPER_MODEL
        self._model = None

    def _load(self):
        if self._model is not None:
            return
        import whisper
        logger.info(f"Loading Whisper '{self.model_name}' model...")
        self._model = whisper.load_model(self.model_name)
        logger.info("Whisper model loaded")

    def transcribe(self, audio_path: str | Path) -> List[Dict]:
        """
        Transcribe audio file.

        Returns:
            List of segment dicts: [{start, end, text}, ...]
        """
        self._load()
        audio_path = str(audio_path)
        logger.info(f"Transcribing: {audio_path}")

        result = self._model.transcribe(
            audio_path,
            language=None,       # auto-detect language
            word_timestamps=False,
            verbose=False,
        )

        segments = [
            {"start": s["start"], "end": s["end"], "text": s["text"].strip()}
            for s in result.get("segments", [])
            if s["text"].strip()
        ]

        full_text = result.get("text", "").strip()
        logger.info(f"Transcribed {len(segments)} segments, {len(full_text)} chars")
        return segments

    def detect_chorus(
        self,
        segments: List[Dict],
        min_similarity: float = 0.65,
        window_size: int = 3,
    ) -> List[float]:
        """
        Find timestamps likely to be chorus (repeated lyric sections).

        Strategy: group segments into rolling windows of `window_size` lines,
        compare every pair with SequenceMatcher. Windows whose text appears
        more than once (similarity >= min_similarity) are flagged as chorus.

        Returns:
            List of start timestamps for detected chorus windows.
        """
        if len(segments) < window_size * 2:
            return []

        # Build text windows
        windows = []
        for i in range(len(segments) - window_size + 1):
            text = " ".join(s["text"] for s in segments[i: i + window_size]).lower()
            start = segments[i]["start"]
            windows.append((start, text))

        chorus_times: List[float] = []

        for i, (t_i, txt_i) in enumerate(windows):
            for j, (t_j, txt_j) in enumerate(windows):
                if j <= i:
                    continue
                # Skip windows that are too close together (same section)
                if abs(t_i - t_j) < 20.0:
                    continue
                ratio = difflib.SequenceMatcher(None, txt_i, txt_j).ratio()
                if ratio >= min_similarity:
                    if t_i not in chorus_times:
                        chorus_times.append(t_i)
                    if t_j not in chorus_times:
                        chorus_times.append(t_j)

        logger.info(f"Detected {len(chorus_times)} chorus timestamps")
        return sorted(chorus_times)

    def full_text(self, segments: List[Dict]) -> str:
        """Join all segment texts into a single string."""
        return " ".join(s["text"] for s in segments).strip()


__all__ = ["Transcriber"]