File size: 491 Bytes
eafbe80 | 1 2 3 4 5 6 7 8 9 10 11 12 13 |
from transformers import AutoConfig, AutoModel, AutoModelForCausalLM
from fla.models.retnet.configuration_retnet import RetNetConfig
from fla.models.retnet.modeling_retnet import RetNetForCausalLM, RetNetModel
AutoConfig.register(RetNetConfig.model_type, RetNetConfig, exist_ok=True)
AutoModel.register(RetNetConfig, RetNetModel, exist_ok=True)
AutoModelForCausalLM.register(RetNetConfig, RetNetForCausalLM, exist_ok=True)
__all__ = ['RetNetConfig', 'RetNetForCausalLM', 'RetNetModel']
|