stt_models.py

"""Motores de reconocimiento de voz de subtify.

Cada motor expone `transcribe(audio_file, language)` y devuelve siempre la misma
estructura, de forma que el resto del pipeline no necesita saber cuál está debajo.

Parakeet es el motor principal: entre 3 y 5 veces más rápido que Whisper y con
mucha menos memoria. Whisper queda como respaldo para los idiomas que Parakeet no
cubre, y como detector de idioma cuando no se especifica ninguno.
"""

import torch

# Los 25 idiomas europeos que cubre Parakeet TDT v3
PARAKEET_LANGUAGES = {
    "en", "es", "fr", "de", "bg", "hr", "cs", "da", "nl", "et", "fi", "el",
    "hu", "it", "lv", "lt", "mt", "pl", "pt", "ro", "sk", "sl", "sv", "ru", "uk",
}

PARAKEET_MODEL = "nvidia/parakeet-tdt-0.6b-v3"

# El respaldo se activa justo en los idiomas difíciles, que es donde large-v3 saca
# más ventaja a la variante turbo. Como serán una minoría, el coste extra compensa.
WHISPER_MODEL = "openai/whisper-large-v3"

# Con chunks de 30 s, un batch de 24 procesa 12 minutos de audio a la vez y dispara
# la memoria por encima de 24 GB. Con 8 el tiempo apenas cambia.
WHISPER_BATCH_SIZE = 8

# Todos los modelos de voz esperan 16 kHz mono
SAMPLE_RATE = 16000


def load_audio(audio_file, max_seconds=None):
    """Carga audio a 16 kHz mono.

    Se usa librosa y no torchaudio: en el Space, torchaudio delega la decodificación
    en torchcodec, que revienta al abrir los mp3 que genera el pipeline. librosa
    hace falta de todas formas, porque el feature extractor de Parakeet la usa.
    """
    import librosa

    audio, _ = librosa.load(
        audio_file, sr=SAMPLE_RATE, mono=True, duration=max_seconds
    )
    return audio


def select_engine(language):
    """Decide qué motor transcribe, según el idioma de origen.

    Args:
        language: código de idioma ('es', 'en', …) o None para autodetectar.

    Returns:
        "parakeet" o "whisper".
    """
    if language is None:
        return "whisper"
    return "parakeet" if language in PARAKEET_LANGUAGES else "whisper"


def detect_language(audio_file, device):
    """Identifica el idioma del audio con Whisper.

    Solo mira los primeros 30 segundos, así que no hace falta transcribir el vídeo
    entero para decidir qué motor usarlo después.
    """
    from transformers import WhisperForConditionalGeneration, WhisperProcessor

    torch_dtype = torch.float16 if device.type == "cuda" else torch.float32

    processor = WhisperProcessor.from_pretrained(WHISPER_MODEL)
    model = WhisperForConditionalGeneration.from_pretrained(
        WHISPER_MODEL, torch_dtype=torch_dtype, low_cpu_mem_usage=True
    ).to(device)

    audio = load_audio(audio_file, max_seconds=30)
    inputs = processor(audio, sampling_rate=SAMPLE_RATE, return_tensors="pt")
    features = inputs.input_features.to(device, dtype=torch_dtype)

    with torch.no_grad():
        detected = model.detect_language(features)

    # detect_language devuelve tokens del tipo <|es|>
    token = processor.tokenizer.convert_ids_to_tokens(detected[0].item())
    return token.strip("<|>")


def transcribe_parakeet(audio_file, device):
    """Transcribe con Parakeet TDT v3, devolviendo timestamps de palabra."""
    try:
        from transformers import AutoModelForTDT, AutoProcessor
    except ImportError as exc:
        # AutoModelForTDT llegó con transformers 5.x. Sin él Parakeet no puede
        # cargarse, y conviene que el error diga qué versión hay instalada.
        import transformers
        raise ImportError(
            f"Parakeet necesita AutoModelForTDT (transformers 5.x). "
            f"Instalada: {transformers.__version__}. Detalle: {exc}"
        ) from exc

    processor = AutoProcessor.from_pretrained(PARAKEET_MODEL)
    model = AutoModelForTDT.from_pretrained(
        PARAKEET_MODEL, dtype="auto", device_map=device
    )

    audio = load_audio(audio_file)
    inputs = processor([audio], sampling_rate=SAMPLE_RATE)
    inputs = inputs.to(model.device, dtype=model.dtype)

    with torch.no_grad():
        output = model.generate(**inputs, return_dict_in_generate=True)

    decoded, timestamps = processor.decode(
        output.sequences, durations=output.durations, skip_special_tokens=True
    )
    text = decoded[0] if isinstance(decoded, (list, tuple)) else decoded
    entries = timestamps[0] if timestamps and isinstance(timestamps[0], (list, tuple)) \
        else timestamps

    subtokens = []
    for entry in entries or []:
        token = entry.get("token", "") if isinstance(entry, dict) else ""
        if not token.strip():
            continue
        subtokens.append((float(entry["start"]), float(entry["end"]), token))

    text = str(text).strip()
    return {"text": text, "chunks": merge_subtokens(subtokens, text)}


def merge_subtokens(subtokens, text):
    """Agrupa los subtokens BPE de Parakeet en palabras.

    Parakeet cronometra subtokens ('Lo', 'ok', 'ing'), no palabras, y un subtítulo
    partido por sílabas no sirve. Como los subtokens concatenados reproducen el
    texto sin espacios, se recorre el texto palabra a palabra consumiendo
    subtokens: el inicio lo marca el primero y el final, el último.
    """
    chunks = []
    idx = 0          # subtoken en curso
    consumed = 0     # caracteres ya usados del subtoken en curso

    for word in text.split():
        target = len(word)
        got = 0
        start = end = None

        while idx < len(subtokens) and got < target:
            s, e, tok = subtokens[idx]
            tok = tok.strip()
            available = len(tok) - consumed
            if available <= 0:
                idx += 1
                consumed = 0
                continue

            if start is None:
                start = s
            end = e

            take = min(available, target - got)
            got += take
            consumed += take

            if consumed >= len(tok):
                idx += 1
                consumed = 0

        if start is not None and end is not None:
            chunks.append({"start": start, "end": end, "text": word})

    return chunks


def transcribe_whisper(audio_file, language, device, chunk_length_s=30,
                       stride_length_s=5):
    """Transcribe con Whisper large-v3, devolviendo timestamps de palabra."""
    from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline

    torch_dtype = torch.float16 if device.type == "cuda" else torch.float32

    model = AutoModelForSpeechSeq2Seq.from_pretrained(
        WHISPER_MODEL,
        torch_dtype=torch_dtype,
        low_cpu_mem_usage=True,
        use_safetensors=True,
        attn_implementation="sdpa",
    ).to(device)
    processor = AutoProcessor.from_pretrained(WHISPER_MODEL)

    pipe = pipeline(
        "automatic-speech-recognition",
        model=model,
        tokenizer=processor.tokenizer,
        feature_extractor=processor.feature_extractor,
        torch_dtype=torch_dtype,
        device=device,
        chunk_length_s=chunk_length_s,
        stride_length_s=stride_length_s,
    )

    result = pipe(
        audio_file,
        return_timestamps="word",
        batch_size=WHISPER_BATCH_SIZE,
        generate_kwargs={
            "language": language,
            "task": "transcribe",
            "use_cache": True,
            "num_beams": 1,
        },
    )

    chunks = []
    for chunk in result.get("chunks", []):
        timestamp = chunk.get("timestamp") or (None, None)
        if timestamp[0] is None or timestamp[1] is None:
            continue
        chunks.append({
            "start": float(timestamp[0]),
            "end": float(timestamp[1]),
            "text": chunk["text"].strip(),
        })

    return {"text": result["text"].strip(), "chunks": chunks}