saracandu
/

stldec_random

Model card Files Files and versions

saracandu commited on Jul 31

Commit

0b21222

·

verified ·

1 Parent(s): 3b2a37f

ultimo tentativo

Files changed (1) hide show

__init__.py +7 -4

__init__.py CHANGED Viewed

@@ -1,8 +1,11 @@
-from transformers import AutoTokenizer, AutoModelForCausalLM, AutoConfig
 from configuration import STLConfig
-from modeling import STLForCausalLM, STLTokenizer
 AutoConfig.register("stldec", STLConfig)
 AutoModelForCausalLM.register(STLConfig, STLForCausalLM)
-from transformers import AutoTokenizer
-AutoTokenizer.register(STLConfig, STLTokenizer)

+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)
+AutoConfig.register_for_auto_class()
+STLModel.register_for_auto_class("AutoModel")
+STLForCausalLM.register_for_auto_class("AutoModelForCausalLM")