Nattapong Tapachoom
Refactor model registration to use CONFIG_MAPPING and MODEL_FOR_CAUSAL_LM_MAPPING
16f8406 | """ | |
| ZombitThai-GPT Model Package | |
| """ | |
| from .modeling import ( | |
| ZombitThaiGPTConfig, | |
| ZombitThaiGPTModel, | |
| ZombitThaiGPTLMHeadModel, | |
| ZombitThaiGPTPreTrainedModel, | |
| ) | |
| # Register the model with transformers | |
| from transformers import AutoConfig, AutoModelForCausalLM | |
| from transformers.models.auto.configuration_auto import CONFIG_MAPPING | |
| from transformers.models.auto.modeling_auto import MODEL_FOR_CAUSAL_LM_MAPPING | |
| # Register the configuration | |
| CONFIG_MAPPING.register("ZombitThaiGPT", ZombitThaiGPTConfig) | |
| # Register the model for causal LM | |
| MODEL_FOR_CAUSAL_LM_MAPPING.register(ZombitThaiGPTConfig, ZombitThaiGPTLMHeadModel) |