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.sse.configuration_sse import SSEConfig
from fla.models.sse.modeling_sse import SSEForCausalLM, SSEModel
AutoConfig.register(SSEConfig.model_type, SSEConfig, exist_ok=True)
AutoModel.register(SSEConfig, SSEModel, exist_ok=True)
AutoModelForCausalLM.register(SSEConfig, SSEForCausalLM, exist_ok=True)
__all__ = ['SSEConfig', 'SSEForCausalLM', 'SSEModel']
|