File size: 449 Bytes
0b21222
92a2d05
0b21222
01cbb24
 
0b21222
01cbb24
0b21222
bdedfcc
0b21222
 
1
2
3
4
5
6
7
8
9
10
11
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")