|
|
import torch |
|
|
import torch.nn as nn |
|
|
import torch.nn.functional as F |
|
|
import numpy as np |
|
|
import pandas as pd |
|
|
from torch.nn.functional import silu |
|
|
from torch.nn.functional import softplus |
|
|
from einops import rearrange, repeat, einsum |
|
|
from transformers import AutoTokenizer, AutoModel |
|
|
from torch import Tensor |
|
|
from einops import rearrange |
|
|
|
|
|
class Embedding(): |
|
|
def __init__(self, model_name='jina', pooling=None): |
|
|
self.model_name = model_name |
|
|
self.pooling = pooling |
|
|
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
|
|
if model_name == 'jina': |
|
|
self.tokenizer = AutoTokenizer.from_pretrained("jinaai/jina-embeddings-v3", code_revision='da863dd04a4e5dce6814c6625adfba87b83838aa', trust_remote_code=True) |
|
|
self.model = AutoModel.from_pretrained("jinaai/jina-embeddings-v3", code_revision='da863dd04a4e5dce6814c6625adfba87b83838aa', trust_remote_code=True).to(self.device) |
|
|
elif model_name == 'xlm-roberta-base': |
|
|
self.tokenizer = AutoTokenizer.from_pretrained('xlm-roberta-base') |
|
|
self.model = AutoModel.from_pretrained('xlm-roberta-base').to(self.device) |
|
|
elif model_name == 'canine-c': |
|
|
self.tokenizer = AutoTokenizer.from_pretrained('google/canine-c') |
|
|
self.model = AutoModel.from_pretrained('google/canine-c').to(self.device) |
|
|
else: |
|
|
raise ValueError('Unknown name of Embedding') |
|
|
def _mean_pooling(self, X): |
|
|
def mean_pooling(model_output, attention_mask): |
|
|
token_embeddings = model_output[0] |
|
|
input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float() |
|
|
return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9) |
|
|
encoded_input = self.tokenizer(X, padding=True, truncation=True, return_tensors='pt').to(self.device) |
|
|
with torch.no_grad(): |
|
|
model_output = self.model(**encoded_input) |
|
|
sentence_embeddings = mean_pooling(model_output, encoded_input['attention_mask']) |
|
|
sentence_embeddings = F.normalize(sentence_embeddings, p=2, dim=1) |
|
|
return sentence_embeddings.unsqueeze(1) |
|
|
|
|
|
def get_embeddings(self, X): |
|
|
if self.pooling is None: |
|
|
if self.model_name == 'canine-c_emb': |
|
|
max_len = 329 |
|
|
else: |
|
|
max_len = 95 |
|
|
encoded_input = self.tokenizer(X, padding=True, truncation=True, return_tensors='pt').to(self.device) |
|
|
with torch.no_grad(): |
|
|
features = self.model(**encoded_input)[0].detach().cpu().float().numpy() |
|
|
res = np.pad(features[:, :max_len, :], ((0, 0), (0, max(0, max_len - features.shape[1])), (0, 0)), "constant") |
|
|
return torch.tensor(res) |
|
|
elif self.pooling == 'mean': |
|
|
return self._mean_pooling(X) |
|
|
else: |
|
|
raise ValueError('Unknown type of pooling') |
|
|
class RMSNorm(nn.Module): |
|
|
def __init__(self, d_model: int, eps: float = 1e-8) -> None: |
|
|
super().__init__() |
|
|
self.eps = eps |
|
|
self.weight = nn.Parameter(torch.ones(d_model)) |
|
|
|
|
|
def forward(self, x: Tensor) -> Tensor: |
|
|
return x * torch.rsqrt(x.pow(2).mean(-1, keepdim = True) + self.eps) * self.weight |
|
|
|
|
|
class Mamba(nn.Module): |
|
|
def __init__(self, num_layers, d_input, d_model, d_state=16, d_discr=None, ker_size=4, num_classes=7, model_name='jina', pooling=None): |
|
|
super().__init__() |
|
|
mamba_par = { |
|
|
'd_input' : d_input, |
|
|
'd_model' : d_model, |
|
|
'd_state' : d_state, |
|
|
'd_discr' : d_discr, |
|
|
'ker_size': ker_size |
|
|
} |
|
|
self.model_name = model_name |
|
|
embed = Embedding(model_name, pooling) |
|
|
self.embedding = embed.get_embeddings |
|
|
self.layers = nn.ModuleList([nn.ModuleList([MambaBlock(**mamba_par), RMSNorm(d_input)]) for _ in range(num_layers)]) |
|
|
self.fc_out = nn.Linear(d_input, num_classes) |
|
|
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
|
|
self.softmax = nn.Softmax(dim=1) |
|
|
|
|
|
def forward(self, seq, cache=None): |
|
|
seq = torch.tensor(self.embedding(seq)).to(self.device) |
|
|
for mamba, norm in self.layers: |
|
|
out, cache = mamba(norm(seq), cache) |
|
|
seq = out + seq |
|
|
return self.fc_out(seq.mean(dim = 1)) |
|
|
|
|
|
def predict(self, x): |
|
|
label_to_emotion = { |
|
|
0: 'anger', |
|
|
1: 'disgust', |
|
|
2: 'fear', |
|
|
3: 'joy/happiness', |
|
|
4: 'neutral', |
|
|
5: 'sadness', |
|
|
6: 'surprise/enthusiasm' |
|
|
} |
|
|
with torch.no_grad(): |
|
|
output = self.forward(x) |
|
|
_, predictions = torch.max(output, dim=1) |
|
|
result = [label_to_emotion[i] for i in (map(int, predictions))] |
|
|
return result |
|
|
|
|
|
def predict_proba(self, x): |
|
|
with torch.no_grad(): |
|
|
output = self.forward(x) |
|
|
|
|
|
return self.softmax(output) |
|
|
|
|
|
class MambaBlock(nn.Module): |
|
|
def __init__(self, d_input, d_model, d_state=16, d_discr=None, ker_size=4): |
|
|
super().__init__() |
|
|
d_discr = d_discr if d_discr is not None else d_model // 16 |
|
|
self.in_proj = nn.Linear(d_input, 2 * d_model, bias=False) |
|
|
self.out_proj = nn.Linear(d_model, d_input, bias=False) |
|
|
self.s_B = nn.Linear(d_model, d_state, bias=False) |
|
|
self.s_C = nn.Linear(d_model, d_state, bias=False) |
|
|
self.s_D = nn.Sequential(nn.Linear(d_model, d_discr, bias=False), nn.Linear(d_discr, d_model, bias=False),) |
|
|
self.conv = nn.Conv1d( |
|
|
in_channels=d_model, |
|
|
out_channels=d_model, |
|
|
kernel_size=ker_size, |
|
|
padding=ker_size - 1, |
|
|
groups=d_model, |
|
|
bias=True, |
|
|
) |
|
|
self.A = nn.Parameter(torch.arange(1, d_state + 1, dtype=torch.float).repeat(d_model, 1)) |
|
|
self.D = nn.Parameter(torch.ones(d_model, dtype=torch.float)) |
|
|
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
|
|
|
|
|
def forward(self, seq, cache=None): |
|
|
b, l, d = seq.shape |
|
|
(prev_hid, prev_inp) = cache if cache is not None else (None, None) |
|
|
a, b = self.in_proj(seq).chunk(2, dim=-1) |
|
|
x = rearrange(a, 'b l d -> b d l') |
|
|
x = x if prev_inp is None else torch.cat((prev_inp, x), dim=-1) |
|
|
a = self.conv(x)[..., :l] |
|
|
a = rearrange(a, 'b d l -> b l d') |
|
|
a = silu(a) |
|
|
a, hid = self.ssm(a, prev_hid=prev_hid) |
|
|
b = silu(b) |
|
|
out = a * b |
|
|
out = self.out_proj(out) |
|
|
if cache: |
|
|
cache = (hid.squeeze(), x[..., 1:]) |
|
|
return out, cache |
|
|
|
|
|
def ssm(self, seq, prev_hid): |
|
|
A = -self.A |
|
|
D = +self.D |
|
|
B = self.s_B(seq) |
|
|
C = self.s_C(seq) |
|
|
s = softplus(D + self.s_D(seq)) |
|
|
A_bar = einsum(torch.exp(A), s, 'd s, b l d -> b l d s') |
|
|
B_bar = einsum( B, s, 'b l s, b l d -> b l d s') |
|
|
X_bar = einsum(B_bar, seq, 'b l d s, b l d -> b l d s') |
|
|
hid = self._hid_states(A_bar, X_bar, prev_hid=prev_hid) |
|
|
out = einsum(hid, C, 'b l d s, b l s -> b l d') |
|
|
out = out + D * seq |
|
|
return out, hid |
|
|
|
|
|
def _hid_states(self, A, X, prev_hid=None): |
|
|
b, l, d, s = A.shape |
|
|
A = rearrange(A, 'b l d s -> l b d s') |
|
|
X = rearrange(X, 'b l d s -> l b d s') |
|
|
if prev_hid is not None: |
|
|
return rearrange(A * prev_hid + X, 'l b d s -> b l d s') |
|
|
h = torch.zeros(b, d, s, device=self.device) |
|
|
return torch.stack([h := A_t * h + X_t for A_t, X_t in zip(A, X)], dim=1) |
|
|
|
