backend/app/features/ssl_extractor.py

"""XLS-R SSL feature extractor (HuggingFace wav2vec2-xls-r-300m).

Used as a shared frontend by Nes2Net and SONAR architectures. Heavy: ~1.2 GB
download on first call. Loaded lazily and only when ``ENABLE_HEAVY_MODELS=true``.
"""
from __future__ import annotations

import threading
from typing import Optional

import torch

from app.logging_setup import get_logger

logger = get_logger(__name__)


class XLSRExtractor:
    """Lazy-loaded XLS-R feature extractor."""

    HF_ID = "facebook/wav2vec2-xls-r-300m"
    HIDDEN_DIM = 1024

    _instance: Optional["XLSRExtractor"] = None
    _lock = threading.Lock()

    def __init__(self, device: Optional[str] = None) -> None:
        from transformers import Wav2Vec2FeatureExtractor, Wav2Vec2Model  # local import keeps import-time light

        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
        logger.info("Loading XLS-R from HuggingFace (%s) on %s ...", self.HF_ID, self.device)
        self.feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(self.HF_ID)
        self.model = Wav2Vec2Model.from_pretrained(self.HF_ID).to(self.device).eval()
        logger.info("XLS-R loaded.")

    @classmethod
    def get(cls) -> "XLSRExtractor":
        with cls._lock:
            if cls._instance is None:
                cls._instance = cls()
            return cls._instance

    @torch.inference_mode()
    def extract(self, waveform: torch.Tensor, sample_rate: int = 16000) -> torch.Tensor:
        """Returns hidden states of shape [1, T_frames, 1024]."""
        if waveform.dim() == 2:
            wav = waveform.squeeze(0).cpu().numpy()
        else:
            wav = waveform.cpu().numpy()
        inputs = self.feature_extractor(wav, sampling_rate=sample_rate, return_tensors="pt")
        input_values = inputs["input_values"].to(self.device)
        out = self.model(input_values)
        return out.last_hidden_state.detach()