File size: 987 Bytes
3275441 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 | """Xavante - audio_encoder.py - Encoder de audio (estilo Whisper leve)."""
from __future__ import annotations
import logging
import torch
import torch.nn as nn
logger = logging.getLogger(__name__)
class AudioEncoder(nn.Module):
"""Convolucoes 1D + Linear para audio bruto ou mel-spec."""
def __init__(self, d_model: int = 512, in_channels: int = 1, n_mels: int = 80):
super().__init__()
in_ch = in_channels if in_channels > 1 else n_mels
self.conv = nn.Sequential(
nn.Conv1d(in_ch, 64, 5, 2, 2),
nn.GELU(),
nn.Conv1d(64, 128, 5, 2, 2),
nn.GELU(),
nn.Conv1d(128, 256, 5, 2, 2),
nn.GELU(),
nn.AdaptiveAvgPool1d(1),
)
self.proj = nn.Linear(256, d_model)
def forward(self, x: torch.Tensor) -> torch.Tensor:
# x: [B, C, T] ou [B, n_mels, T]
feat = self.conv(x).squeeze(-1)
return self.proj(feat)
__all__ = ["AudioEncoder"]
|