|
|
| 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 |
|
|
