from transformers import PreTrainedModel, PretrainedConfig
from .GPTModel import GPTModel

class MyGPTConfig(PretrainedConfig):
    model_type = "my_gpt"
    def __init__(self, vocab_size=50257, context_length=256, emb_dim=128,
                 n_layers=2, n_heads=2, drop_rate=0.1, **kwargs):
        super().__init__(**kwargs)
        self.vocab_size = vocab_size
        self.context_length = context_length
        self.emb_dim = emb_dim
        self.n_layers = n_layers
        self.n_heads = n_heads
        self.drop_rate = drop_rate

class MyGPTModelForCausalLM(PreTrainedModel):
    config_class = MyGPTConfig

    def __init__(self, config):
        super().__init__(config)
        self.inner = GPTModel(config.to_dict())   # 用你的 GPTModel

    def forward(self, input_ids, **kwargs):
        logits = self.inner(input_ids)
        return {"logits": logits}