# Copyright 2026 Roblox Corporation
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
Self-contained inference for VoiceToxicityClassifier models saved as safetensors + config.json.

Dependencies: torch, safetensors, numpy. Optional: transformers (for AutoModel.from_pretrained).

Model directory must contain:
  - config.json (hidden_size, num_labels, labels, etc.)
  - model.safetensors

Usage:
  # As module
  from inference import load_model, load_audio, run_inference, load_audio_batch
  from inference import extract_label_scores
  model, config = load_model("/path/to/model_dir")
  audio = load_audio("clip.wav", config=config)
  raw = run_inference(model, audio)
  out = extract_label_scores(raw["probs"], raw.get("language_probs"), config, index=0)
  # Batch: multiple WAVs
  batch_audio, mask = load_audio_batch(["a.wav", "b.wav"], config=config)
  raw = run_inference(model, batch_audio, attention_mask=mask)

  # CLI (single or multiple WAVs; model_dir defaults to current directory)
  python inference.py clip.wav
  python inference.py --model-dir /path/to/model_dir clip.wav
  python inference.py a.wav b.wav c.wav --output results.json
"""

import argparse
import json
import math
import wave
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple, Union

import numpy as np
import torch
from safetensors.torch import load_file
from torch import nn

try:
    from transformers import AutoConfig, AutoModel
    from transformers.configuration_utils import PretrainedConfig
    from transformers.modeling_utils import PreTrainedModel

    _HF_AVAILABLE = True
except ImportError:
    _HF_AVAILABLE = False
    AutoConfig = None  # type: ignore[misc, assignment]
    AutoModel = None  # type: ignore[misc, assignment]
    PreTrainedConfig = None  # type: ignore[misc, assignment]
    PreTrainedModel = nn.Module  # type: ignore[misc, assignment]

EXPECTED_SAMPLE_RATE = 16000


def max_audio_samples_from_config(cfg: Dict[str, Any]) -> int:
    """Raw audio sample cap: ``max_positions * hop_length * 2`` (matches model truncation)."""
    return int(cfg["max_positions"]) * int(cfg["hop_length"]) * 2


class VoiceToxicityClassifierConfig(PretrainedConfig if _HF_AVAILABLE else object):
    model_type = "VoiceToxicityClassifier"

    def __init__(
        self,
        n_mels: int = 80,
        n_fft: int = 400,
        hop_length: int = 160,
        dropout: float = 0.0,
        hidden_size: int = 1024,
        num_attention_heads: int = 16,
        num_hidden_layers: int = 24,
        classifier_proj_size: int = 640,
        max_positions: int = 1500,
        num_labels: int = 2,
        num_language_heads: int = 0,
        time_reduction: Optional[List[List[int]]] = None,
        **kwargs: Any,
    ):
        if _HF_AVAILABLE:
            super().__init__(**kwargs)
        self.n_mels = n_mels
        self.n_fft = n_fft
        self.hop_length = hop_length
        self.dropout = dropout
        self.hidden_size = hidden_size
        self.num_attention_heads = num_attention_heads
        self.num_hidden_layers = num_hidden_layers
        self.classifier_proj_size = classifier_proj_size
        self.max_positions = max_positions
        self.num_labels = num_labels
        self.num_language_heads = num_language_heads
        self.time_reduction = time_reduction or [[4, 2]]


if _HF_AVAILABLE:
    AutoConfig.register("VoiceToxicityClassifier", VoiceToxicityClassifierConfig)


@dataclass
class _PoolingOutput:
    logits: torch.Tensor
    language_logits: Optional[torch.Tensor]


class _LogMelFeatureExtraction(nn.Module):
    """Log-mel spectrogram from raw audio. Weights (fb, window) loaded from state_dict."""

    def __init__(self, n_mels: int = 80, n_fft: int = 400, hop_length: int = 160):
        super().__init__()
        self.n_fft = n_fft
        self.n_mels = n_mels
        self.hop_length = hop_length
        # Shapes: fb (n_mels, n_fft//2+1), window (n_fft)
        self.register_buffer("fb", torch.zeros(n_mels, n_fft // 2 + 1))
        self.register_buffer("window", torch.hann_window(n_fft))

    def forward(self, wav: torch.Tensor) -> torch.Tensor:
        stft = torch.stft(
            wav, self.n_fft, self.hop_length, window=self.window, return_complex=True
        )
        stft = torch.view_as_real(stft)
        real, imag = stft[..., 0], stft[..., 1]
        real, imag = real[..., :-1], imag[..., :-1]
        magnitudes = real**2 + imag**2
        mel_spec = self.fb @ magnitudes
        log_spec = torch.clamp(mel_spec, min=1e-10).log10()
        log_spec = torch.maximum(log_spec, log_spec.max() - 8.0)
        return (log_spec + 4.0) / 4.0


class _ConvolutionalReductionModule(nn.Module):
    def __init__(
        self, input_dim: int, projection_dim: int, kernel_size: int, stride: int
    ):
        super().__init__()
        self.stride = stride
        self.pre_conv_seq = nn.Sequential(
            nn.Linear(input_dim, 2 * projection_dim, bias=True),
            nn.GLU(dim=2),
        )
        self.conv_module = nn.Conv1d(
            projection_dim,
            projection_dim,
            kernel_size,
            stride=stride,
            dilation=1,
            padding=(kernel_size - 1) // 2,
            groups=projection_dim,
            bias=True,
        )
        self.post_conv_seq = nn.Sequential(
            nn.LayerNorm(projection_dim),
            nn.SiLU(),
            nn.Linear(projection_dim, projection_dim, bias=True),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.pre_conv_seq(x)
        x = self.conv_module(x.transpose(1, 2)).transpose(1, 2)
        return self.post_conv_seq(x)


class _SelfAttentionPooling(nn.Module):
    def __init__(self, input_dim: int):
        super().__init__()
        self.W = nn.Linear(input_dim, 1)

    def forward(
        self, inp: torch.Tensor, attention_mask: Optional[torch.Tensor] = None
    ) -> torch.Tensor:
        att_w = self.W(inp).squeeze(-1)
        if attention_mask is not None:
            att_w = att_w.masked_fill(attention_mask == 0, float("-inf"))
        att_w = torch.softmax(att_w, dim=1).unsqueeze(-1)
        return torch.sum(inp * att_w, dim=1)


class _TorchAttentionPoolingClassifier(nn.Module):
    def __init__(
        self,
        hidden_size: int,
        classifier_proj_size: int,
        num_labels: int,
        language_heads: int = 0,
    ):
        super().__init__()
        self.conv_reduction1 = _ConvolutionalReductionModule(
            hidden_size, hidden_size, 7, 2
        )
        self.mid_layer_norm1 = nn.LayerNorm(hidden_size)
        self.mid_attention1 = nn.MultiheadAttention(
            embed_dim=hidden_size, num_heads=16, batch_first=True
        )
        self.pre_conv2_layer_norm = nn.LayerNorm(hidden_size)
        self.conv_reduction2 = _ConvolutionalReductionModule(
            hidden_size, classifier_proj_size, 7, 2
        )
        self.mid_layer_norm2 = nn.LayerNorm(classifier_proj_size)
        self.mid_attention2 = nn.MultiheadAttention(
            embed_dim=classifier_proj_size, num_heads=16, batch_first=True
        )
        self.pre_conv3_layer_norm = nn.LayerNorm(classifier_proj_size)
        self.conv_reduction3 = _ConvolutionalReductionModule(
            classifier_proj_size, classifier_proj_size, 5, 2
        )
        self.pre_pooling_layer_norm = nn.LayerNorm(classifier_proj_size)
        self.pooling_attention = _SelfAttentionPooling(classifier_proj_size)
        self.classifier = nn.Linear(classifier_proj_size, num_labels)

        self.language_heads = None
        self.language_projector = None
        if language_heads > 0:
            self.language_projector = nn.Linear(hidden_size, classifier_proj_size)
            self.language_heads = nn.Linear(classifier_proj_size, language_heads)

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
    ) -> _PoolingOutput:
        original_input = hidden_states
        orig_attention_mask = attention_mask
        hidden_states = self.conv_reduction1(hidden_states)
        if attention_mask is not None:
            attention_mask = attention_mask[:, :: self.conv_reduction1.stride]

        residual = hidden_states
        hidden_states = self.mid_layer_norm1(hidden_states)
        key_padding_mask = None if attention_mask is None else ~attention_mask
        hidden_states, _ = self.mid_attention1(
            hidden_states,
            hidden_states,
            hidden_states,
            key_padding_mask=key_padding_mask,
        )
        hidden_states = hidden_states + residual
        hidden_states = self.pre_conv2_layer_norm(hidden_states)
        hidden_states = self.conv_reduction2(hidden_states)
        if attention_mask is not None:
            attention_mask = attention_mask[:, :: self.conv_reduction2.stride]

        residual = hidden_states
        hidden_states = self.mid_layer_norm2(hidden_states)
        key_padding_mask = None if attention_mask is None else ~attention_mask
        hidden_states, _ = self.mid_attention2(
            hidden_states,
            hidden_states,
            hidden_states,
            key_padding_mask=key_padding_mask,
        )
        hidden_states = hidden_states + residual
        hidden_states = self.pre_conv3_layer_norm(hidden_states)
        hidden_states = self.conv_reduction3(hidden_states)
        if attention_mask is not None:
            attention_mask = attention_mask[:, :: self.conv_reduction3.stride]

        hidden_states = self.pre_pooling_layer_norm(hidden_states)
        pooled = self.pooling_attention(hidden_states, attention_mask)
        logits = self.classifier(pooled)

        language_logits = None
        if self.language_heads is not None:
            proj = self.language_projector(original_input)
            if orig_attention_mask is not None:
                lang_mask = orig_attention_mask[:, : proj.shape[1]]
                denom = torch.sum(orig_attention_mask, dim=1, keepdim=True)
                pooled_lang = torch.sum(proj * lang_mask.unsqueeze(-1), dim=1) / denom
            else:
                pooled_lang = proj.mean(dim=1)
            language_logits = self.language_heads(pooled_lang)

        return _PoolingOutput(
            logits=logits,
            language_logits=language_logits,
        )


def _sinusoids(length: int, channels: int, max_timescale: float = 2000) -> torch.Tensor:
    if channels % 2 != 0:
        raise ValueError("channels must be even for sinusoids")
    log_inc = math.log(max_timescale) / (channels // 2 - 1)
    inv_timescales = torch.exp(-log_inc * torch.arange(channels // 2))
    scaled = torch.arange(length).view(-1, 1) * inv_timescales.view(1, -1)
    return torch.cat([scaled.sin(), scaled.cos()], dim=1)


class VoiceToxicityClassifier(
    PreTrainedModel if _HF_AVAILABLE else nn.Module  # type: ignore[misc]
):
    """Inference-only VoiceToxicityClassifier.
    When transformers is installed, supports AutoModel.from_pretrained(model_dir).
    """

    config_class = VoiceToxicityClassifierConfig

    def __init__(self, config: VoiceToxicityClassifierConfig):
        if _HF_AVAILABLE:
            super().__init__(config)
        else:
            super().__init__()
        self.hop_length = config.hop_length
        self.max_samples = config.max_positions * config.hop_length * 2
        self.time_reduction = [
            (int(layer), int(ratio)) for layer, ratio in config.time_reduction
        ]
        self.time_reduction_pooling = nn.ModuleDict()
        for layer_idx, ratio in self.time_reduction:
            self.time_reduction_pooling[str(layer_idx)] = nn.AvgPool1d(ratio)

        self.feature_extractor = _LogMelFeatureExtraction(
            config.n_mels, config.n_fft, config.hop_length
        )
        self.conv1 = nn.Conv1d(
            config.n_mels, config.hidden_size, kernel_size=3, padding=1
        )
        self.conv2 = nn.Conv1d(
            config.hidden_size, config.hidden_size, kernel_size=3, stride=2, padding=1
        )
        self.embed_positions = nn.Embedding(config.max_positions, config.hidden_size)
        self.attn_mask_conv = nn.Conv1d(1, 1, kernel_size=3, stride=2, padding=1)
        self.dropout = config.dropout

        self.layers = nn.ModuleList()
        for _ in range(config.num_hidden_layers):
            self.layers.append(
                nn.TransformerEncoderLayer(
                    d_model=config.hidden_size,
                    nhead=config.num_attention_heads,
                    dim_feedforward=config.hidden_size * 4,
                    dropout=config.dropout,
                    activation="gelu",
                    batch_first=True,
                    norm_first=True,
                    bias=True,
                )
            )
        self.final_layer_norm = nn.LayerNorm(config.hidden_size)

        self.classifier = _TorchAttentionPoolingClassifier(
            config.hidden_size,
            config.classifier_proj_size,
            config.num_labels,
            config.num_language_heads,
        )

        # Fixed init for positional embedding and mask conv (loaded from state_dict)
        with torch.no_grad():
            self.embed_positions.weight.copy_(
                _sinusoids(config.max_positions, config.hidden_size)
            )
            nn.init.zeros_(self.attn_mask_conv.weight)
            nn.init.zeros_(self.attn_mask_conv.bias)
            self.attn_mask_conv.weight.data[:, :, 2] = torch.eye(1, 1)
        self.embed_positions.requires_grad_(False)
        for p in self.attn_mask_conv.parameters():
            p.requires_grad = False

        if _HF_AVAILABLE:
            self.post_init()

    def forward(
        self,
        input_values: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        return_dict: bool = True,
    ) -> Union[Dict[str, Optional[torch.Tensor]], torch.Tensor]:
        input_values = input_values[:, : self.max_samples]
        input_features = self.feature_extractor(input_values).to(
            self.embed_positions.weight.dtype
        )
        if attention_mask is not None:
            attention_mask = attention_mask[:, : self.max_samples]
            attention_mask = attention_mask[
                :, : -(self.hop_length - 1) : self.hop_length
            ]

        inputs_embeds = torch.nn.functional.gelu(self.conv1(input_features))
        inputs_embeds = torch.nn.functional.gelu(self.conv2(inputs_embeds))

        if attention_mask is not None:
            attention_mask = self.attn_mask_conv(
                attention_mask.unsqueeze(1).float()
            ).squeeze(1)
            attention_mask = attention_mask != 0

        inputs_embeds = inputs_embeds.permute(0, 2, 1)
        embed_pos = self.embed_positions.weight
        hidden_states = inputs_embeds + embed_pos[: inputs_embeds.shape[1], :]
        hidden_states = torch.nn.functional.dropout(
            hidden_states, p=self.dropout, training=self.training
        )

        # TransformerEncoderLayer expects True = pad/ignore; mask is 1 = valid.
        layer_attention_mask = None if attention_mask is None else (attention_mask == 0)
        tr_pooling = self.time_reduction_pooling

        for idx, encoder_layer in enumerate(self.layers):
            hidden_states = encoder_layer(
                hidden_states,
                src_key_padding_mask=layer_attention_mask,
            )
            idx_str = str(idx)
            if idx_str in tr_pooling:
                pool = tr_pooling[idx_str]
                hidden_states = pool(hidden_states.transpose(1, 2)).transpose(1, 2)
                ratio = pool.kernel_size[0]
                if attention_mask is not None:
                    attention_mask = attention_mask[:, ::ratio]
                    attention_mask = attention_mask[:, : hidden_states.shape[1]]
                    # TransformerEncoderLayer expects True = pad/ignore; mask is 1 = valid.
                    layer_attention_mask = attention_mask == 0

        hidden_states = self.final_layer_norm(hidden_states)
        if attention_mask is not None:
            attention_mask = attention_mask[:, : hidden_states.shape[1]]
        pooling_outputs = self.classifier(hidden_states, attention_mask)

        if return_dict:
            return {
                "logits": pooling_outputs.logits,
                "language_logits": pooling_outputs.language_logits,
            }
        if pooling_outputs.language_logits is not None:
            return torch.cat(
                [pooling_outputs.logits, pooling_outputs.language_logits], dim=1
            )
        return pooling_outputs.logits


if _HF_AVAILABLE:
    AutoModel.register(VoiceToxicityClassifierConfig, VoiceToxicityClassifier)  # type: ignore[attr-defined]


def load_model(
    model_dir: Union[str, Path],
    device: Optional[Union[str, torch.device]] = None,
) -> Tuple[nn.Module, Dict[str, Any]]:
    """
    Load model and config from a directory with config.json and model.safetensors.
    Returns (model, config). Model is in eval mode.

    When transformers is installed, you can instead use:
      from transformers import AutoModel
      model = AutoModel.from_pretrained("/path/to/model_dir")
    For that, config.json should include "model_type": "VoiceToxicityClassifier".
    """
    model_dir = Path(model_dir)
    with open(model_dir / "config.json", encoding="utf-8") as f:
        cfg = json.load(f)
    if "model_type" not in cfg:
        cfg["model_type"] = "VoiceToxicityClassifier"
    config = VoiceToxicityClassifierConfig(**cfg)
    model = VoiceToxicityClassifier(config)
    state = load_file(model_dir / "model.safetensors")
    missing, unexpected = model.load_state_dict(state, strict=False)
    if missing:
        raise RuntimeError(f"Missing keys when loading state_dict: {missing}")
    if unexpected:
        raise RuntimeError(f"Unexpected keys in state_dict: {unexpected}")

    model.eval()
    if device is not None:
        model = model.to(device)
    return model, cfg


def load_audio(
    path: Union[str, Path],
    target_sr: int = EXPECTED_SAMPLE_RATE,
    max_length: Optional[int] = None,
    config: Optional[Dict[str, Any]] = None,
) -> torch.Tensor:
    """
    Load a WAV file to a float32 tensor, shape (1, samples).
    Expects mono or takes first channel. Sample rate must match target_sr.
    If ``max_length`` is None, ``config`` must be provided; the cap is
    ``max_audio_samples_from_config(config)``, matching ``load_audio_batch`` and model
    truncation. Longer audio is truncated to the first ``max_length`` samples.
    """
    path = Path(path)
    with wave.open(str(path), "rb") as wav:
        nch = wav.getnchannels()
        sr = wav.getframerate()
        nframes = wav.getnframes()
        raw = wav.readframes(nframes)
    dtype = np.int16
    samples = np.frombuffer(raw, dtype=dtype)
    if nch > 1:
        samples = samples.reshape(-1, nch)[:, 0]
    audio = samples.astype(np.float32) / 32768.0

    if sr != target_sr:
        raise RuntimeError(
            f"Audio sample rate {sr} Hz does not match expected {target_sr} Hz: {path}"
        )
    if max_length is None:
        if config is None:
            max_length = 480000  # 480000 samples @ 16kHz = 30s
        else:
            max_length = max_audio_samples_from_config(config)
    if audio.shape[0] > max_length:
        audio = audio[:max_length]
    return torch.from_numpy(audio).unsqueeze(0)


def load_audio_batch(
    paths: List[Union[str, Path]],
    target_sr: int = EXPECTED_SAMPLE_RATE,
    max_length: Optional[int] = None,
    config: Optional[Dict[str, Any]] = None,
) -> Tuple[torch.Tensor, torch.Tensor]:
    """
    Load multiple WAV files into a batched tensor and an attention mask.

    Returns:
        batch_audio: (B, max_len) float32 tensor, right-padded with zeros.
        attention_mask: (B, max_len) 1 for valid samples, 0 for padding.
    """
    if max_length is None:
        if config is None:
            max_length = 480000  # 480000 samples @ 16kHz = 30s
        else:
            max_length = max_audio_samples_from_config(config)
    audios = []
    lengths = []
    for path in paths:
        a = load_audio(path, target_sr=target_sr, max_length=max_length)
        audios.append(a.squeeze(0))
        lengths.append(a.shape[1])
    max_len = min(max(lengths), max_length)
    batch = torch.zeros(len(audios), max_len, dtype=torch.float32)
    mask = torch.zeros(len(audios), max_len, dtype=torch.long)
    for i, (a, L) in enumerate(zip(audios, lengths)):
        L = min(L, max_len)
        batch[i, :L] = a[:L]
        mask[i, :L] = 1
    return batch, mask


def run_inference(
    model: nn.Module,
    audio: torch.Tensor,
    attention_mask: Optional[torch.Tensor] = None,
    device: Optional[Union[str, torch.device]] = None,
) -> Dict[str, Any]:
    """
    Run inference on audio tensor(s).

    If attention_mask is None, treats audio as single (shape [1, samples] or [samples])
    and uses a mask of ones. Otherwise treats as batch (shape [B, samples]) with the
    given attention mask.

    Returns dict with logits, probs, and optionally language_logits, language_probs (all on device).
    """
    if device is None:
        device = next(model.parameters()).device
    if attention_mask is None:
        if audio.dim() == 1:
            audio = audio.unsqueeze(0)
        attention_mask = torch.ones(
            audio.shape[0], audio.shape[1], dtype=torch.long, device=device
        )
    audio = audio.to(device)
    attention_mask = attention_mask.to(device)
    with torch.no_grad():
        out = model(audio, attention_mask=attention_mask, return_dict=True)
    logits = out["logits"].float()
    language_logits = out.get("language_logits")
    if language_logits is not None:
        language_logits = language_logits.float()
    probs = torch.sigmoid(logits)
    language_probs = (
        torch.softmax(language_logits, dim=-1) if language_logits is not None else None
    )
    return {
        "logits": logits,
        "probs": probs,
        "language_logits": language_logits,
        "language_probs": language_probs,
    }


def extract_label_scores(
    probs: torch.Tensor,
    language_probs: Optional[torch.Tensor],
    config: Dict[str, Any],
    index: Optional[int] = None,
) -> Dict[str, Any]:
    """
    Extract label_scores and optionally language_probs dicts from raw probability tensors.
    For batch tensors, pass index to select one row.
    """
    if index is not None:
        probs = probs[index]
        if language_probs is not None:
            language_probs = language_probs[index]
    probs_np = probs.cpu().numpy()
    labels = config.get("labels", [])
    result: Dict[str, Any] = {
        "label_scores": {
            (labels[i] if i < len(labels) else str(i)): float(probs_np[i])
            for i in range(probs.shape[-1])
        }
    }
    if language_probs is not None:
        lang_probs_np = language_probs.cpu().numpy()
        languages = config.get("languages", [])
        result["language_probs"] = {
            (languages[i] if i < len(languages) else str(i)): float(lang_probs_np[i])
            for i in range(language_probs.shape[-1])
        }
    return result


def _write_single_result(
    output_path: str, wav_path: str, result: Dict[str, Any]
) -> None:
    out = {
        "path": str(Path(wav_path).resolve()),
        "label_scores": result["label_scores"],
        "language_probs": result.get("language_probs"),
    }
    with open(output_path, "w", encoding="utf-8") as f:
        json.dump(out, f, indent=2)
    print("Wrote", output_path)


def _write_batch_results(
    output_path: str,
    wav_paths: List[str],
    results: List[Dict[str, Any]],
) -> None:
    out = {
        "results": [
            {
                "path": str(Path(p).resolve()),
                "label_scores": r["label_scores"],
                "language_probs": r.get("language_probs"),
            }
            for p, r in zip(wav_paths, results)
        ]
    }
    with open(output_path, "w", encoding="utf-8") as f:
        json.dump(out, f, indent=2)
    print("Wrote", output_path)


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Run VoiceToxicityClassifier inference on WAV file(s)"
    )
    parser.add_argument(
        "--model-dir",
        "-m",
        default=".",
        help="Directory with config.json and model.safetensors (default: current directory)",
    )
    parser.add_argument(
        "wav_paths",
        nargs="+",
        help="One or more WAV files at 16 kHz (mono or first channel of stereo)",
    )
    parser.add_argument(
        "--device",
        default="cpu",
        help="Device: cpu or cuda (default: cpu)",
    )
    parser.add_argument(
        "--output",
        default=None,
        help="Optional JSON path to write results (for batch: list of results keyed by path)",
    )
    args = parser.parse_args()
    device = args.device if torch.cuda.is_available() or args.device == "cpu" else "cpu"
    model, config = load_model(args.model_dir, device=device)

    if len(args.wav_paths) == 1:
        audio = load_audio(args.wav_paths[0], config=config)
        raw = run_inference(model, audio, device=device)
        display = extract_label_scores(
            raw["probs"], raw.get("language_probs"), config, index=0
        )
        print("Label scores:", display["label_scores"])
        if "language_probs" in display:
            print("Language probs:", display["language_probs"])
        if args.output:
            _write_single_result(args.output, args.wav_paths[0], display)
    else:
        batch_audio, mask = load_audio_batch(args.wav_paths, config=config)
        raw = run_inference(model, batch_audio, attention_mask=mask, device=device)
        batch_size = raw["logits"].shape[0]
        results = [
            extract_label_scores(
                raw["probs"], raw.get("language_probs"), config, index=i
            )
            for i in range(batch_size)
        ]
        for path, res in zip(args.wav_paths, results):
            print(path, "->", res["label_scores"])
            if "language_probs" in res:
                print("  language_probs:", res["language_probs"])
        if args.output:
            _write_batch_results(args.output, args.wav_paths, results)


if __name__ == "__main__":
    main()