from transformers import PretrainedConfig, PreTrainedModel
from transformer_engine.pytorch import TransformerLayer


class SimpleTEConfig(PretrainedConfig):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.hidden_size = 320
        self.intermediate_size = 1024
        self.num_attention_heads = 16


class SimpleTEModel(PreTrainedModel):
    config_class = SimpleTEConfig

    def __init__(self, config: SimpleTEConfig):
        super().__init__(config)
        self.te_layer = TransformerLayer(
            hidden_size=config.hidden_size,
            ffn_hidden_size=config.intermediate_size,
            num_attention_heads=config.num_attention_heads,
        )

    def forward(self, hidden_states, attention_mask):
        return self.te_layer(hidden_states, attention_mask)


def test_simple_te_model(tmp_path):
    config = SimpleTEConfig()
    model = SimpleTEModel(config)

    model.save_pretrained(tmp_path / "simple_te_model")
    del model
    model = SimpleTEModel.from_pretrained(tmp_path / "simple_te_model")
    assert isinstance(model.te_layer, TransformerLayer)