File size: 507 Bytes
eafbe80 | 1 2 3 4 5 6 7 8 9 10 11 12 13 |
from transformers import AutoConfig, AutoModel, AutoModelForCausalLM
from fla.models.rodimus.configuration_rodimus import RodimusConfig
from fla.models.rodimus.modeling_rodimus import RodimusForCausalLM, RodimusModel
AutoConfig.register(RodimusConfig.model_type, RodimusConfig, exist_ok=True)
AutoModel.register(RodimusConfig, RodimusModel, exist_ok=True)
AutoModelForCausalLM.register(RodimusConfig, RodimusForCausalLM, exist_ok=True)
__all__ = ['RodimusConfig', 'RodimusForCausalLM', 'RodimusModel']
|