File size: 2,332 Bytes
44a3c2c | 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 |
import torch
import torch.nn as nn
from transformers import PretrainedConfig, PreTrainedModel
class SupernovaEncoderConfig(PretrainedConfig):
model_type = 'supernova_encoder'
def __init__(
self,
vocab_size=50257,
hidden_size=512,
num_hidden_layers=6,
num_attention_heads=8,
intermediate_size=2048,
max_position_embeddings=300,
output_dim=2304,
layer_norm_eps=1e-12,
**kwargs
):
super().__init__(**kwargs)
self.vocab_size = vocab_size
self.hidden_size = hidden_size
self.num_hidden_layers = num_hidden_layers
self.num_attention_heads = num_attention_heads
self.intermediate_size = intermediate_size
self.max_position_embeddings = max_position_embeddings
self.output_dim = output_dim
self.layer_norm_eps = layer_norm_eps
class SupernovaNepaliEncoder(PreTrainedModel):
config_class = SupernovaEncoderConfig
def __init__(self, config):
super().__init__(config)
self.embeddings = nn.Embedding(config.vocab_size, config.hidden_size)
self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size)
layer = nn.TransformerEncoderLayer(
d_model=config.hidden_size,
nhead=config.num_attention_heads,
dim_feedforward=config.intermediate_size,
batch_first=True,
norm_first=True
)
self.encoder = nn.TransformerEncoder(layer, num_layers=config.num_hidden_layers)
self.projection = nn.Linear(config.hidden_size, config.output_dim)
self.ln_final = nn.LayerNorm(config.output_dim, eps=config.layer_norm_eps)
self.post_init()
def forward(self, input_ids, attention_mask=None):
seq_length = input_ids.size(1)
position_ids = torch.arange(seq_length, dtype=torch.long, device=input_ids.device).unsqueeze(0)
x = self.embeddings(input_ids) + self.position_embeddings(position_ids)
padding_mask = None
if attention_mask is not None:
padding_mask = ~(attention_mask.bool())
hidden_states = self.encoder(x, src_key_padding_mask=padding_mask)
projected = self.projection(hidden_states)
return projected
|