| # -*- coding: utf-8 -*- | |
| from transformers import AutoConfig, AutoModel, AutoModelForCausalLM | |
| from fla.models.transformer.configuration_transformer import TransformerConfig | |
| from fla.models.transformer.modeling_transformer import ( | |
| TransformerForCausalLM, TransformerModel) | |
| AutoConfig.register(TransformerConfig.model_type, TransformerConfig) | |
| AutoModel.register(TransformerConfig, TransformerModel) | |
| AutoModelForCausalLM.register(TransformerConfig, TransformerForCausalLM) | |
| __all__ = ['TransformerConfig', 'TransformerForCausalLM', 'TransformerModel'] | |