Upload 4 files

text/architectures.py

@@ -0,0 +1,202 @@
+# coding: utf-8
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+from .help_layers import TransformerEncoderLayer, CustomMambaBlock
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+
+class EmotionMamba(nn.Module):
+    def __init__(self, input_dim_emotion=1024, input_dim_personality=1024, hidden_dim=128, out_features=512, mamba_layer_number=2, positional_encoding=True, num_transformer_heads=4, transformer_dropout=0.1, tr_layer_number=1, dropout=0.1, num_emotions=7, num_traits=5):
+        super().__init__()
+
+        self.hidden_dim = hidden_dim
+
+        self.emo_proj = nn.Sequential(
+            nn.Linear(input_dim_emotion, hidden_dim),
+            nn.LayerNorm(hidden_dim),
+            nn.Dropout(dropout)
+        )
+
+        self.emotion_encoder = nn.ModuleList([
+            CustomMambaBlock(hidden_dim, hidden_dim, dropout=dropout)
+            for _ in range(mamba_layer_number)
+        ])
+
+
+        self.emotion_fc_out = nn.Sequential(
+            nn.Linear(hidden_dim, out_features),
+            nn.LayerNorm(out_features),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(out_features, num_emotions)
+        )
+
+
+    def forward(self, emotion_input=None, personality_input=None, return_features=False):
+        emo = self.emo_proj(emotion_input)  # (B, T, hidden_dim)
+        for layer in self.emotion_encoder:
+            emo = layer(emo)
+        out_emo = self.emotion_fc_out(emo.mean(dim=1))  # (B, num_emotions)
+        if return_features:
+            return {
+                'emotion_logits': out_emo,
+                'last_encoder_features': emo,
+            }
+        else:
+            return {'emotion_logits': out_emo}
+
+
+class PersonalityMamba(nn.Module):
+    def __init__(self, input_dim_emotion=1024, input_dim_personality=1024, hidden_dim=128, out_features=512, mamba_layer_number=2, per_activation="sigmoid", positional_encoding=True, num_transformer_heads=4, tr_layer_number=1, dropout=0.1, num_emotions=7, num_traits=5):
+        super().__init__()
+
+        self.hidden_dim = hidden_dim
+
+        self.per_proj = nn.Sequential(
+            nn.Linear(input_dim_personality, hidden_dim),
+            nn.LayerNorm(hidden_dim),
+            nn.Dropout(dropout)
+        )
+
+        self.personality_encoder = nn.ModuleList([
+            CustomMambaBlock(hidden_dim, hidden_dim, dropout=dropout)
+            for _ in range(mamba_layer_number)
+        ])
+
+        self.personality_fc_out = nn.Sequential(
+            nn.Linear(hidden_dim, out_features),
+            nn.LayerNorm(out_features),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(out_features, num_traits)
+        )
+
+        if per_activation == "sigmoid":
+            self.activation = nn.Sigmoid()
+        elif per_activation == "relu":
+            self.activation = nn.ReLU()
+
+    def forward(self, emotion_input=None, personality_input=None, return_features=False):
+        per = self.per_proj(personality_input)
+
+        for layer in self.personality_encoder:
+            per = layer(per)
+
+        out_per = self.personality_fc_out(per.mean(dim=1))
+    
+        if return_features:
+            return {
+                'personality_scores': self.activation(out_per),
+                'last_encoder_features': per,
+            }
+        else:
+            return {'personality_scores': self.activation(out_per)}
+
+
+class FusionTransformer(nn.Module):
+    def __init__(self, emo_model, per_model, hidden_dim=128, out_features=512, per_activation="sigmoid", positional_encoding=True, num_transformer_heads=4, tr_layer_number=1, dropout=0.1, num_emotions=7, num_traits=5):
+        super().__init__()
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+
+        self.hidden_dim = hidden_dim
+
+        self.emo_model = emo_model
+        self.per_model = per_model
+
+        for param in self.emo_model.parameters():
+            param.requires_grad = False
+
+        for param in self.per_model.parameters():
+            param.requires_grad = False
+
+        self.emo_proj = nn.Sequential(
+            nn.Linear(self.emo_model.hidden_dim, hidden_dim),
+            nn.LayerNorm(hidden_dim),
+            nn.Dropout(dropout)
+        )
+
+        self.per_proj = nn.Sequential(
+            nn.Linear(self.per_model.hidden_dim, hidden_dim),
+            nn.LayerNorm(hidden_dim),
+            nn.Dropout(dropout)
+        )
+
+        self.emotion_to_personality_attn = nn.ModuleList([
+            TransformerEncoderLayer(
+                input_dim=hidden_dim,
+                num_heads=num_transformer_heads,
+                dropout=dropout,
+                positional_encoding=positional_encoding
+            ) for _ in range(tr_layer_number)
+        ])
+
+        self.personality_to_emotion_attn = nn.ModuleList([
+            TransformerEncoderLayer(
+                input_dim=hidden_dim,
+                num_heads=num_transformer_heads,
+                dropout=dropout,
+                positional_encoding=positional_encoding
+            ) for _ in range(tr_layer_number)
+        ])
+
+        self.emotion_personality_fc_out = nn.Sequential(
+            nn.Linear(hidden_dim*2, out_features),
+            nn.LayerNorm(out_features),
+            nn.SiLU(),
+            nn.Dropout(dropout),
+            nn.Linear(out_features, num_emotions)
+        )
+
+        self.personality_emotion_fc_out = nn.Sequential(
+            nn.Linear(hidden_dim*2, out_features),
+            nn.LayerNorm(out_features),
+            nn.SiLU(),
+            nn.Dropout(dropout),
+            nn.Linear(out_features, num_traits)
+        )        
+
+        if per_activation == "sigmoid":
+            self.activation = nn.Sigmoid()
+        elif per_activation == "relu":
+            self.activation = nn.ReLU()
+
+    def forward(self, emotion_input=None, personality_input=None, return_features=False):
+        emo_features = self.emo_model(emotion_input=emotion_input, return_features=True)
+        per_features = self.per_model(personality_input=personality_input, return_features=True)
+
+        emo_emd = self.emo_proj(emo_features['last_encoder_features'])
+        per_emd = self.per_proj(per_features['last_encoder_features'])
+  
+        # padding
+        max_len = max(emo_emd.shape[1], per_emd.shape[1])
+        emo_emd = emo_emd.cpu().detach().numpy()
+        per_emd = per_emd.cpu().detach().numpy()
+        emo_emd = np.pad(emo_emd[:, :max_len, :], ((0, 0), (0, max(0, max_len - emo_emd.shape[1])), (0, 0)), "constant")
+        per_emd = np.pad(per_emd[:, :max_len, :], ((0, 0), (0, max(0, max_len - per_emd.shape[1])), (0, 0)), "constant")
+        emo_emd = torch.tensor(emo_emd, device=self.device)
+        per_emd = torch.tensor(per_emd, device=self.device)
+
+        for layer in self.emotion_to_personality_attn:
+            emo_emd += layer(emo_emd, per_emd, per_emd)
+
+        for layer in self.personality_to_emotion_attn:
+            per_emd += layer(per_emd, emo_emd, emo_emd)
+
+        fused = torch.cat([emo_emd, per_emd], dim=-1)
+        emotion_logits = self.emotion_personality_fc_out(fused.mean(dim=1))
+        personality_scores = self.personality_emotion_fc_out(fused.mean(dim=1))
+
+        if return_features:
+            return {
+                'emotion_logits': (emotion_logits+emo_features['emotion_logits'])/2,
+                'personality_scores': (self.activation(personality_scores)+per_features['personality_scores'])/2,
+                'last_emo_encoder_features': emo_emd,
+                'last_per_encoder_features': per_emd,
+            }
+        else:
+            return {'emotion_logits': (emotion_logits+emo_features['emotion_logits'])/2,
+                    'personality_scores': (self.activation(personality_scores)+per_features['personality_scores'])/2,}

text/feature_extractor.py

@@ -0,0 +1,51 @@
+# coding: utf-8
+import torch
+from transformers import AutoTokenizer, AutoModel
+from .model_loader import load_fusion_model
+
+
+class PretrainedTextEmbeddingExtractor:
+    """
+    jinaai/jina-embeddings-v → последовательный эмбеддинг (B, T, 1024) →
+    Fusion-модель → логиты эмоций, оценки Big-5 и последние признаки.
+    """
+
+    def __init__(
+        self,
+        device: str = "cuda",
+        model_name: str = "jinaai/jina-embeddings-v3",
+        fusion_ckpt: str = "modalities/text/checkpoints_models/Transformer_jina_fusion.pt",
+        emo_ckpt: str   = "modalities/text/checkpoints_models/Mamba_jina_emotion.pt",
+        per_ckpt: str   = "modalities/text/checkpoints_models/Mamba_jina_personality.pt",
+    ):
+        self.device = torch.device(device)
+
+        self.tok = AutoTokenizer.from_pretrained(model_name, code_revision='da863dd04a4e5dce6814c6625adfba87b83838aa', trust_remote_code=True)
+        self.enc = AutoModel.from_pretrained(model_name, code_revision='da863dd04a4e5dce6814c6625adfba87b83838aa', trust_remote_code=True).to(self.device).eval()
+
+        self.fusion, _ = load_fusion_model(
+            fusion_ckpt, emo_ckpt, per_ckpt, device=self.device
+        )
+
+    @torch.no_grad()
+    def extract(self, texts: list[str] | str) -> dict:
+        if isinstance(texts, str):
+            texts = [texts]
+
+        batch = self.tok(texts, padding=True, truncation=True,
+                        return_tensors="pt").to(self.device)
+
+        hidden = self.enc(**batch).last_hidden_state  # (B, T, 1024)
+
+        out = self.fusion(
+            emotion_input=hidden.float(),
+            personality_input=hidden.float(),
+            return_features=True,
+        )
+
+        return {
+            "emotion_logits": out["emotion_logits"].cpu(),
+            "personality_scores": out["personality_scores"].cpu(),
+            "last_emo_encoder_features": out["last_emo_encoder_features"].cpu(),
+            "last_per_encoder_features": out["last_per_encoder_features"].cpu(),
+        }

text/help_layers.py

@@ -0,0 +1,98 @@
+# coding: utf-8
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.nn.init as init
+import numpy as np
+import math
+
+
+class PositionWiseFeedForward(nn.Module):
+    def __init__(self, input_dim, hidden_dim, dropout=0.1):
+        super().__init__()
+        self.layer_1 = nn.Linear(input_dim, hidden_dim)
+        self.layer_2 = nn.Linear(hidden_dim, input_dim)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x):
+        x = self.layer_1(x)
+        x = F.gelu(x)
+        x = self.dropout(x)
+        return self.layer_2(x)
+
+
+class AddAndNorm(nn.Module):
+    def __init__(self, input_dim, dropout=0.1):
+        super().__init__()
+        self.norm = nn.LayerNorm(input_dim)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x, residual):
+        return self.norm(x + self.dropout(residual))
+
+
+class PositionalEncoding(nn.Module):
+    def __init__(self, d_model, dropout=0.1, max_len=5000):
+        super().__init__()
+        self.dropout = nn.Dropout(p=dropout)
+
+        position = torch.arange(max_len).unsqueeze(1)
+        div_term = torch.exp(torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model))
+        pe = torch.zeros(max_len, d_model)
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+
+        self.register_buffer("pe", pe)
+
+    def forward(self, x):
+        x = x + self.pe[: x.size(1)].detach()  # Отключаем градиенты
+        return self.dropout(x)
+
+
+class TransformerEncoderLayer(nn.Module):
+    def __init__(self, input_dim, num_heads, dropout=0.1, positional_encoding=False):
+        super().__init__()
+        self.input_dim = input_dim
+        self.self_attention = nn.MultiheadAttention(input_dim, num_heads, dropout=dropout, batch_first=True)
+        self.feed_forward = PositionWiseFeedForward(input_dim, input_dim, dropout=dropout)
+        self.add_norm_after_attention = AddAndNorm(input_dim, dropout=dropout)
+        self.add_norm_after_ff = AddAndNorm(input_dim, dropout=dropout)
+        self.positional_encoding = PositionalEncoding(input_dim) if positional_encoding else None
+
+    def forward(self, query, key, value):
+        if self.positional_encoding:
+            key = self.positional_encoding(key)
+            value = self.positional_encoding(value)
+            query = self.positional_encoding(query)
+
+        attn_output, _ = self.self_attention(query, key, value, need_weights=False)
+
+        x = self.add_norm_after_attention(attn_output, query)
+
+        ff_output = self.feed_forward(x)
+        x = self.add_norm_after_ff(ff_output, x)
+
+        return x
+
+class CustomMambaBlock(nn.Module):
+    def __init__(self, d_input, d_model, dropout=0.1):
+        super().__init__()
+        self.in_proj = nn.Linear(d_input, d_model)
+        self.s_B = nn.Linear(d_model, d_model)
+        self.s_C = nn.Linear(d_model, d_model)
+        self.out_proj = nn.Linear(d_model, d_input)
+        self.norm = nn.LayerNorm(d_input)
+        self.dropout = nn.Dropout(dropout)
+        self.activation = nn.ReLU()
+
+    def forward(self, x):
+        x_in = x
+        x = self.in_proj(x)
+        B = self.s_B(x)
+        C = self.s_C(x)
+        x = x + B + C
+        x = self.activation(x)
+        x = self.out_proj(x)
+        x = self.dropout(x)
+        x = self.norm(x + x_in)
+        return x

text/model_loader.py

@@ -0,0 +1,80 @@
+# coding: utf-8
+import torch
+from .architectures import (
+    EmotionMamba,
+    PersonalityMamba,
+    FusionTransformer,
+)
+
+
+def load_pretrained_emotion_encoder(checkpoint_path, device):
+    emotion_model = EmotionMamba(
+        input_dim_emotion=1024,
+        input_dim_personality=1024,
+        hidden_dim=256,
+        out_features=128,
+        mamba_layer_number=2,
+        dropout=0.1
+    ).to(device)
+
+    checkpoint = torch.load(checkpoint_path, map_location=device)
+    state_dict = checkpoint["model_state_dict"] if "model_state_dict" in checkpoint else checkpoint
+    emotion_model.load_state_dict(state_dict)
+
+    def extract_features(inputs, lengths):
+        features = emotion_model.emo_proj(inputs)
+        for block in emotion_model.emotion_encoder:
+            features = block(features)
+        return features
+
+    emotion_model.extract_features = extract_features
+    emotion_model.eval()
+    return emotion_model
+
+def load_pretrained_personality_encoder(checkpoint_path, device):
+    personality_model = PersonalityMamba(
+        input_dim_emotion=1024, 
+        input_dim_personality=1024, 
+        hidden_dim=64, 
+        out_features=256, 
+        mamba_layer_number=3, 
+        dropout=0.1).to(device)
+
+    checkpoint = torch.load(checkpoint_path, map_location=device)
+    personality_model.load_state_dict(checkpoint)
+
+    def extract_features(inputs, lengths):
+        features = personality_model.per_proj(inputs)
+        for block in personality_model.personality_encoder:
+            features = block(features, features, features)
+        return features
+
+    personality_model.extract_features = extract_features
+    personality_model.eval()
+    return personality_model
+
+def load_fusion_model(
+    fusion_checkpoint_path: str,
+    emotion_encoder_checkpoint: str,
+    personality_encoder_checkpoint: str,
+    device: str = "cpu",
+):
+    device = torch.device(device)
+
+    emotion_encoder = load_pretrained_emotion_encoder(emotion_encoder_checkpoint, device)
+    personality_encoder = load_pretrained_personality_encoder(personality_encoder_checkpoint, device)
+
+    checkpoint = torch.load(fusion_checkpoint_path, map_location=device)
+
+    fusion_model = FusionTransformer(
+        emo_model=emotion_encoder,
+        per_model=personality_encoder,
+        hidden_dim=128,
+        out_features=64,
+        tr_layer_number=3,
+        num_transformer_heads=16,
+        dropout=0.1
+    ).to(device)
+    fusion_model.load_state_dict(checkpoint)
+    fusion_model.eval()
+    return fusion_model, device