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") |