File size: 3,397 Bytes
1d2de52 | 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 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 | import torch.nn as nn
import numpy as np
import torch
from transformers import PreTrainedModel
from .config import FeelWiseConfig
class PositionalEncoding(nn.Module):
def __init__(self, d_model, max_len=500):
super(PositionalEncoding, self).__init__()
self.d_model = d_model
pos = np.arange(max_len)[:, np.newaxis]
i = np.arange(d_model)[np.newaxis, :]
angle_rates = 1 / np.power(10000, (2 * (i // 2)) / np.float32(d_model))
pos_encoding = pos * angle_rates
pos_encoding[:, 0::2] = np.sin(pos_encoding[:, 0::2])
pos_encoding[:, 1::2] = np.cos(pos_encoding[:, 1::2])
self.pos_encoding = torch.tensor(pos_encoding, dtype=torch.float32)
def forward(self, x):
x = x * np.sqrt(self.d_model)
x = x + self.pos_encoding[:x.size(1), :].to(x.device)
return x
class AddNorm(nn.Module):
def __init__(self, d_model):
super(AddNorm, self).__init__()
self.layer_norm = nn.LayerNorm(d_model)
def forward(self, x, sub_layer_x):
return self.layer_norm(x + sub_layer_x)
class FeedForward(nn.Module):
def __init__(self, d_model, d_ff):
super(FeedForward, self).__init__()
self.linear1 = nn.Linear(d_model, d_ff)
self.linear2 = nn.Linear(d_ff, d_model)
self.relu = nn.ReLU()
def forward(self, x):
return self.linear2(self.relu(self.linear1(x)))
class EncoderLayer(nn.Module):
def __init__(self, d_model, n_head, d_ff, dropout=0.1):
super(EncoderLayer, self).__init__()
self.multi_head_attention = nn.MultiheadAttention(d_model, n_head)
self.add_norm1 = AddNorm(d_model)
self.feed_forward = FeedForward(d_model, d_ff)
self.add_norm2 = AddNorm(d_model)
self.dropout = nn.Dropout(dropout)
def forward(self, x):
sub_layer_x = self.multi_head_attention(x, x, x)[0]
sub_layer_x = self.dropout(sub_layer_x)
x = self.add_norm1(x, sub_layer_x)
sub_layer_x = self.feed_forward(x)
sub_layer_x = self.dropout(sub_layer_x)
x = self.add_norm2(x, sub_layer_x)
return x
class Encoder(nn.Module):
def __init__(self, n_layers, d_model, max_len, input_vocab_size, n_head, d_ff, dropout=0.1):
super(Encoder, self).__init__()
self.layers = nn.ModuleList([EncoderLayer(d_model, n_head, d_ff, dropout) for _ in range(n_layers)])
self.embedding = nn.Embedding(input_vocab_size, d_model)
self.pos_encoding = PositionalEncoding(d_model, max_len)
self.dropout = nn.Dropout(dropout)
def forward(self, x):
x = self.embedding(x)
x = self.pos_encoding(x)
x = self.dropout(x)
for layer in self.layers:
x = layer(x)
return x
class FeelWiseModel(PreTrainedModel):
config_class = FeelWiseConfig
base_model_prefix = "FeelWiseEmotion"
def __init__(self, config):
super().__init__(config)
self.encoder = Encoder(config.n_layers, config.d_model, config.max_len, config.input_vocab_size, config.n_head, config.d_ff, config.dropout)
self.fc = nn.Linear(config.d_model, config.num_classes) # Final classification layer
def forward(self, input_ids):
x = self.encoder(input_ids) # Include attention_mask if your encoder uses it
x = x.mean(dim=1)
logits = self.fc(x)
return logits |