stldec_random / __init__.py
saracandu's picture
come non detto
bdedfcc verified
raw
history blame contribute delete
449 Bytes
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")