from transformers import AutoConfig, AutoModel, AutoModelForCausalLM
from configuration import STLConfig
from modeling import STLModel, STLForCausalLM, STLTokenizer

AutoConfig.register("stldec", STLConfig)
AutoModel.register(STLConfig, STLModel)
AutoModelForCausalLM.register(STLConfig, STLForCausalLM)

STLConfig.register_for_auto_class()
STLModel.register_for_auto_class("AutoModel")
STLForCausalLM.register_for_auto_class("AutoModelForCausalLM")