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