create spaces

app.py

@@ -0,0 +1,68 @@
+import gradio as gr
+import torch
+import torchaudio
+from hydra.utils import instantiate
+from src.models import AudioBatch
+
+# Путь к чекпоинтам
+CHECKPOINTS = {
+    "RU+EN": "gigaam_ru_en/gigaam_ru_en.ckpt"
+}
+
+# Кэш моделей
+LOADED_MODELS = {}
+
+def load_model(ckpt_path):
+    if ckpt_path in LOADED_MODELS:
+        return LOADED_MODELS[ckpt_path]
+
+    checkpoint = torch.load(ckpt_path, map_location='cpu')
+    config = checkpoint['config']
+    id2name = checkpoint['id2name']
+
+    model = instantiate(config,  _recursive_=False)
+    model.load_state_dict(checkpoint['state_dict'])
+    model.eval()
+
+    LOADED_MODELS[ckpt_path] = (model, id2name)
+    return model, id2name
+
+def classify_emotion(audio, ckpt_name):
+    # Load waveform
+    waveform, sr = torchaudio.load(audio)
+
+    # Load model
+    model, id2name = load_model(CHECKPOINTS[ckpt_name])
+
+    # Если нужно, ресемплим до 16к
+    if sr != 16000:
+        waveform = torchaudio.transforms.Resample(orig_freq=sr, new_freq=16000)(waveform)
+
+    # B x T
+    if waveform.dim() > 1:
+        waveform = waveform.mean(dim=0, keepdim=True)
+
+    length = torch.tensor([waveform.shape[-1]])
+
+    batch = AudioBatch(waveform, length, None)
+
+    with torch.no_grad():
+        logits, _ = model(batch)
+        probs = torch.softmax(logits, dim=-1).squeeze(0)
+
+    result = {label: float(probs[i]) for i, label in enumerate(id2name)}
+    return result
+
+demo = gr.Interface(
+    fn=classify_emotion,
+    inputs=[
+        gr.Audio(type="filepath", label="Загрузите аудиофайл"),
+        gr.Dropdown(choices=list(CHECKPOINTS.keys()), label="Выберите модель", value="RU+EN")
+    ],
+    outputs=gr.Label(label="Эмоциональная окраска (вероятности)"),
+    title="Эмоциональная классификация речи (дообученная GigaAM на 8 классов)",
+    description="Выберите модель и загрузите аудио"
+)
+
+
+demo.launch()

gigaam_ru_en/gigaam_ru_en.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c6508c61dd2a7093161a5a8faa052a4b3f1dfb5916ba68c5701e8968068598d0
+size 968547236

requirements.txt

@@ -0,0 +1,4 @@
+hydra-core
+gigaam
+gradio
+soundfile

src/losses.py

@@ -0,0 +1,292 @@
+import math 
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class AMSoftmax(nn.Module):
+    def __init__(self, in_features, out_features, s=30.0, m=0.35):
+        """
+        in_features: размерность входных эмбеддингов
+        out_features: количество классов
+        s: масштабный множитель (scale)
+        m: аддитивный margin
+        """
+        super(AMSoftmax, self).__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        self.s = s
+        self.m = m
+        self.weight = nn.Parameter(torch.FloatTensor(out_features, in_features))
+        nn.init.xavier_uniform_(self.weight)
+        self.use_labels_when_train = True
+
+    def inference(self, x):
+        x_norm = F.normalize(x, p=2, dim=1)
+        w_norm = F.normalize(self.weight, p=2, dim=1)
+        logits = F.linear(x_norm, w_norm) * self.s
+        return logits
+
+    def forward(self, x, labels=None):
+        if not self.training or labels is None:
+            return self.inference(x)
+        # Нормализация входов и весов
+        input_norm = F.normalize(x, p=2, dim=1)
+        weight_norm = F.normalize(self.weight, p=2, dim=1)
+
+        # Косинус угла между входами и центрами классов
+        cosine = F.linear(input_norm, weight_norm)  # [batch_size, num_classes]
+
+        # Скопировать для дальнейшего вычисления
+        phi = cosine - self.m
+
+        # Создать one-hot маску
+        one_hot = torch.zeros_like(cosine)
+        one_hot.scatter_(1, labels.view(-1, 1), 1.0)
+
+        # Применить margin только к целевым логитам
+        output = self.s * (one_hot * phi + (1.0 - one_hot) * cosine)
+        return output
+
+
+
+class AAMSoftmax(nn.Module):
+    def __init__(self, in_features, out_features, s=30.0, m=0.50, easy_margin=False):
+        """
+        in_features: размерность эмбеддинга
+        out_features: количество классов
+        s: scale (обычно 30)
+        m: angular margin (обычно 0.5 радиан)
+        easy_margin: использовать "easy margin" трюк или нет
+        """
+        super(AAMSoftmax, self).__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        self.s = s
+        self.m = m
+        self.easy_margin = easy_margin
+
+        self.weight = nn.Parameter(torch.FloatTensor(out_features, in_features))
+        nn.init.xavier_uniform_(self.weight)
+
+        self.cos_m = math.cos(m)
+        self.sin_m = math.sin(m)
+        self.th = math.cos(math.pi - m)
+        self.mm = math.sin(math.pi - m) * m
+        self.use_labels_when_train = True
+
+    def inference(self, x):
+        x_norm = F.normalize(x, p=2, dim=1)
+        w_norm = F.normalize(self.weight, p=2, dim=1)
+        logits = F.linear(x_norm, w_norm) * self.s
+        return logits
+
+    def forward(self, x, labels=None):
+        if not self.training or labels is None:
+            return self.inference(x)
+        # Нормализуем входы и веса
+        cosine = F.linear(F.normalize(x), F.normalize(self.weight))  # [B, C]
+        sine = torch.sqrt(1.0 - cosine ** 2 + 1e-6)
+
+        # cos(θ + m) = cosθ * cos(m) - sinθ * sin(m)
+        phi = cosine * self.cos_m - sine * self.sin_m
+
+        if self.easy_margin:
+            # Используем "легкий" трюк, чтобы избежать неустойчивости
+            phi = torch.where(cosine > 0, phi, cosine)
+        else:
+            # Ограничиваем phi снизу
+            phi = torch.where(cosine > self.th, phi, cosine - self.mm)
+
+        # One-hot метки
+        one_hot = torch.zeros_like(cosine)
+        one_hot.scatter_(1, labels.view(-1, 1), 1.0)
+
+        # Вычисляем итоговый логит
+        output = self.s * (one_hot * phi + (1.0 - one_hot) * cosine)
+        return output
+
+
+class RAMSoftmax(nn.Module):
+    """
+    Real Additive Margin Softmax (RAM-Softmax)
+
+    Args:
+        in_features: размерность входных эмбеддингов
+        out_features: число классов
+        s: scale-фактор для логитов
+        m: additive margin
+        eps: небольшой стабилизатор для sqrt
+    """
+
+    def __init__(self, in_features, out_features, s=30.0, m=0.35, eps=1e-6):
+        super(RAMSoftmax, self).__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        self.s = s
+        self.m = m
+        self.eps = eps
+        # веса центров классов
+        self.weight = nn.Parameter(torch.FloatTensor(out_features, in_features))
+        nn.init.xavier_uniform_(self.weight)
+
+        # Большое отрицательное для «отсечки» легко разделённых классов
+        self.register_buffer('large_neg', torch.tensor(-1e9))
+        self.use_labels_when_train = True
+
+    def inference(self, x):
+        x_norm = F.normalize(x, p=2, dim=1)  # [B, D]
+        w_norm = F.normalize(self.weight, p=2, dim=1)  # [C, D]
+
+        # 2) косинус
+        cosine = torch.matmul(x_norm, w_norm.t())  # [B, C]
+
+        # 3) масштаб
+        logits = cosine * self.s  # [B, C]
+        return logits
+
+    def forward(self, x, labels=None):
+        if not self.training or labels is None:
+
+            return self.inference(x)
+        # 1) нормализуем эмбеддинги и веса
+        x_norm = F.normalize(x, p=2, dim=1)  # [B, D]
+        w_norm = F.normalize(self.weight, p=2, dim=1)  # [C, D]
+
+        # 2) вычисляем все косинусы
+        cosine = F.linear(x_norm, w_norm)  # [B, C]
+
+        # 3) достаём целевой косинус и вычитаем margin
+        idx = torch.arange(x.size(0), device=x.device)
+        cos_y = cosine[idx, labels]  # [B]
+        phi = cos_y - self.m  # [B]
+
+        # 4) заменяем целевой логит на phi, остальные — оставляем как cosine
+        logits = cosine.clone()
+        logits[idx, labels] = phi
+
+        # 5) маскирование «легко» разделённых: для каждого j≠y
+        #    если cos_y - cos_j > m  => отсечь (логит -> large_neg)
+        diff = cos_y.unsqueeze(1) - cosine  # [B, C]
+        mask = (diff > self.m)  # [B, C]
+        mask[idx, labels] = False  # не маскируем целевой
+        logits = torch.where(mask, self.large_neg, logits)
+
+        # 6) масштабируем
+        logits = logits * self.s
+
+        return logits
+
+
+class RAAMSoftmax(nn.Module):
+    def __init__(self, in_features, out_features, s=30.0, m=0.50, eps=1e-6):
+        super().__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        self.s = s
+        self.m = m
+        self.eps = eps
+
+        self.weight = nn.Parameter(torch.FloatTensor(out_features, in_features))
+        nn.init.xavier_uniform_(self.weight)
+
+        self.cos_m = math.cos(m)
+        self.sin_m = math.sin(m)
+        self.large_neg = -1e9  # для masking
+
+    def inference(self, x):
+        x_norm = F.normalize(x, p=2, dim=1)
+        w_norm = F.normalize(self.weight, p=2, dim=1)
+        logits = F.linear(x_norm, w_norm) * self.s
+        return logits
+    
+    def forward(self, x, labels=None):
+        if not self.training or labels is None:
+            return self.inference(x)
+        x = F.normalize(x, dim=1)
+        W = F.normalize(self.weight, dim=1)
+
+        cosine = F.linear(x, W)  # [B, C]
+        sine = torch.sqrt(1.0 - cosine ** 2 + self.eps)
+
+        cos_theta_y = cosine[torch.arange(x.size(0)), labels]
+        phi = cos_theta_y * self.cos_m - sine[torch.arange(x.size(0)), labels] * self.sin_m
+
+        logits = cosine.clone()
+        logits[torch.arange(x.size(0)), labels] = phi
+
+        # RAM: отсечка легкоразделённых негативов
+        diff = cos_theta_y.unsqueeze(1) - cosine
+        mask = (diff > self.m)
+        mask[torch.arange(x.size(0)), labels] = False
+
+        logits = torch.where(mask, self.large_neg, logits)
+
+        return self.s * logits
+
+
+class WeightCrossEntropy(nn.Module):
+    def __init__(self, id2name: list, class_distribution: dict):
+        super().__init__()
+        weight = torch.Tensor([1 / math.sqrt(class_distribution[name]) for name in id2name])
+        self.ce = nn.CrossEntropyLoss(weight=weight)
+    def forward(self, input, target):
+        return self.ce(input, target)
+
+
+class CrossEntropyLabelSmooth(nn.Module):
+    """Cross entropy loss with label smoothing regularizer.
+
+    Reference:
+    Szegedy et al. Rethinking the Inception Architecture for Computer Vision. CVPR 2016.
+    Equation: y = (1 - epsilon) * y + epsilon / K.
+
+    Args:
+        num_classes (int): number of classes.
+        epsilon (float): weight.
+    """
+
+    def __init__(self, num_classes, epsilon=0.1, id2name=None, class_distribution=None):
+        super(CrossEntropyLabelSmooth, self).__init__()
+        self.num_classes = num_classes
+        self.epsilon = epsilon
+        self.logsoftmax = nn.LogSoftmax(dim=1)
+        if id2name is not None and class_distribution is not None:
+            weights = torch.Tensor([1 / math.sqrt(class_distribution[name]) for name in id2name])
+            self.weights = weights.to(torch.float32)
+        else:
+            self.weights = None
+
+    def forward(self, inputs, targets, use_label_smoothing=True):
+        """
+        Args:
+            inputs: prediction matrix (before softmax) with shape (batch_size, num_classes)
+            targets: ground truth labels with shape (b,)
+        """
+        #targets = torch.zeros(labels.size(0), self.num_classes).to(labels.device)
+        #targets.scatter_(1, labels.unsqueeze(1), 1)
+        #targets = targets.long()
+        log_probs = self.logsoftmax(inputs)
+        targets = torch.zeros(log_probs.size()).scatter_(1, targets.unsqueeze(1).data.cpu(), 1).to(targets.device)
+        #if self.use_gpu: targets = targets #.to(torch.device('cuda'))
+        if use_label_smoothing:
+            targets = (1 - self.epsilon) * targets + self.epsilon / self.num_classes
+        loss = (- targets * log_probs)
+        if self.weights is not None:
+            weights = self.weights.to(loss.device)
+            loss = loss * weights.unsqueeze(0)  # (batch_size, num_classes)
+        loss = loss.sum(dim=1).mean()
+        return loss
+
+
+
+
+class AMSoftmaxLoss(nn.Module):
+    def __init__(self, in_features, out_features, num_classes, s=30.0, m=0.50, easy_margin=False, epsilon=0.1, id2name=None, class_distribution=None):
+        super(AMSoftmaxLoss, self).__init__()
+        self.aam = AAMSoftmax(in_features, out_features, s, m, easy_margin)
+        self.criterion = CrossEntropyLabelSmooth(num_classes, epsilon, id2name, class_distribution)
+
+    def forward(self, inputs, targets):
+        return self.criterion(self.aam(inputs, targets), targets)

src/models.py

@@ -0,0 +1,123 @@
+
+
+from typing import Dict, List, Tuple, Union
+
+import hydra
+import omegaconf
+import torch
+from torch import Tensor, nn
+
+
+import omegaconf
+from gigaam.model import GigaAM, GigaAMEmo
+from gigaam.preprocess import SAMPLE_RATE, load_audio
+from gigaam.utils import onnx_converter
+
+import torch.nn.functional as F
+from dataclasses import dataclass
+@dataclass
+class AudioBatch:
+    wavs: torch.Tensor
+    wav_lengths: torch.Tensor
+    emotions: torch.LongTensor
+
+class FeatureExtractorGigaAM(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        
+        #checkpoint = torch.load(path_to_weight, map_location="cpu", weights_only=False)
+  
+        self.fe = GigaAM(cfg)
+        #self.fe.load_state_dict(checkpoint["state_dict"], strict=False)
+    def forward(self, features, feature_lengths): #input raw wavs, attention mask [B, WAV_LEN]
+        return self.fe(features, feature_lengths) # return [B, EMB_DIM, T], [B]
+     
+class MaxPooling(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, features, feature_lengths):
+        # features: [B, T, D]
+        features = features.transpose(1, 2)  # → [B, D, T]
+        pooled = F.max_pool1d(features, kernel_size=features.shape[-1])  # → [B, D, 1]
+        return pooled.squeeze(-1)  # → [B, D]
+
+
+class AttentionStatsPooling(nn.Module):
+    def __init__(self, input_dim, attn_dim=128):
+        super().__init__()
+        self.attn = nn.Sequential(
+            nn.Linear(input_dim, attn_dim),
+            nn.Tanh(),
+            nn.Linear(attn_dim, 1)
+        )
+
+    def forward(self, x, lens):  # x: [B, T, D], lens: [B]
+        B, T, D = x.size()
+        device = x.device
+
+        # [B, T, 1]
+        attn_scores = self.attn(x)
+
+        # создаём маску: [B, T]
+        mask = torch.arange(T, device=device).unsqueeze(0) < lens.unsqueeze(1)  # [B, T]
+        mask = mask.unsqueeze(-1)  # [B, T, 1]
+
+        # маскируем паддинг
+        attn_scores[~mask] = float('-inf')
+
+        # softmax по валидным позициям
+        attn_weights = F.softmax(attn_scores, dim=1)  # [B, T, 1]
+        attn_weights = attn_weights * mask  # зануляем padded веса (на всякий случай)
+
+        # считаем взвешенное среднее и std
+        mean = torch.sum(attn_weights * x, dim=1)  # [B, D]
+        std = torch.sqrt(torch.sum(attn_weights * (x - mean.unsqueeze(1))**2, dim=1) + 1e-9)
+
+        return torch.cat([mean, std], dim=1)  # [B, 2*D]
+
+
+class SelfAttentionWithStatsPooling(nn.Module):
+    def __init__(self, embed_dim=768, num_heads=4, attn_dim=128, out_dim=256):
+        super().__init__()
+        self.mha = nn.MultiheadAttention(embed_dim=embed_dim, num_heads=num_heads, batch_first=True)
+        self.norm = nn.LayerNorm(embed_dim)
+        self.pool = AttentionStatsPooling(input_dim=embed_dim, attn_dim=attn_dim)
+        self.out = nn.Linear(2 * embed_dim, out_dim)
+
+    def forward(self, x, lens):  # x: [B, T, D], lens: [B]
+        B, T, D = x.size()
+        device = x.device
+
+        # Attention mask: [B, T]
+        attn_mask = torch.arange(T, device=device).unsqueeze(0) >= lens.unsqueeze(1)  # pad == True
+
+        # Преобразуем для MultiheadAttention: [B, T] → [B, T] → [B, T] → [B, T] (bool)
+
+        attn_out, _ = self.mha(x, x, x, key_padding_mask=attn_mask)  # [B, T, D]
+        attn_out = self.norm(attn_out + x)
+
+        pooled = self.pool(attn_out, lens)  # [B, 2*D]
+        return self.out(pooled)             # [B, out_dim]
+
+
+
+
+class EmotionModel(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+
+        self.feature_extractor = FeatureExtractorGigaAM(config.feature_extractor.cfg)#hydra.utils.instantiate(config.feature_extractor)
+        self.pooling = hydra.utils.instantiate(config.pooling)
+        self.head = hydra.utils.instantiate(config.head)
+    
+    def forward(self, batch: AudioBatch):
+        feats, lengths = self.feature_extractor(batch.wavs, batch.wav_lengths) # return  [B, EMB_DIM, T], [B]
+        feats = feats.transpose(1, 2)# [B, T, EMB_DIM]
+
+        pooled = self.pooling(feats, lengths) # [B, NEW_EMB_DIM]
+        if hasattr(self.head, "use_labels_when_train") and self.head.use_labels_when_train is True:
+            logit = self.head(pooled, batch.emotions)
+        else:
+            logit = self.head(pooled)
+        return logit, None # [B, NUM_CLASSES]