File size: 442 Bytes
eafbe80 | 1 2 3 4 5 6 7 8 9 10 11 12 |
from transformers import AutoConfig, AutoModel, AutoModelForCausalLM
from fla.models.kda.configuration_kda import KDAConfig
from fla.models.kda.modeling_kda import KDAForCausalLM, KDAModel
AutoConfig.register(KDAConfig.model_type, KDAConfig, exist_ok=True)
AutoModel.register(KDAConfig, KDAModel, exist_ok=True)
AutoModelForCausalLM.register(KDAConfig, KDAForCausalLM, exist_ok=True)
__all__ = ['KDAConfig', 'KDAForCausalLM', 'KDAModel']
|