File size: 1,079 Bytes
f45ecb6 7b87608 f45ecb6 7b87608 f45ecb6 7b87608 f45ecb6 7b87608 f45ecb6 |
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 |
# PyPilot Model Architecture
import torch
import torch.nn as nn
from transformers import PreTrainedModel, PretrainedConfig
class PyPilotConfig(PretrainedConfig):
model_type = "pypilot"
def __init__(self, vocab_size=50000, hidden_size=768, num_layers=12, **kwargs):
self.vocab_size = vocab_size
self.hidden_size = hidden_size
self.num_layers = num_layers
super().__init__(**kwargs)
class PyPilotModel(PreTrainedModel):
config_class = PyPilotConfig
def __init__(self, config):
super().__init__(config)
self.embedding = nn.Embedding(config.vocab_size, config.hidden_size)
self.transformer_blocks = nn.ModuleList([
nn.TransformerEncoderLayer(config.hidden_size, 8)
for _ in range(config.num_layers)
])
self.output_layer = nn.Linear(config.hidden_size, config.vocab_size)
def forward(self, input_ids):
x = self.embedding(input_ids)
for block in self.transformer_blocks:
x = block(x)
return self.output_layer(x) |