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