from transformers import AutoConfig, AutoModel, AutoModelForCausalLM, AutoTokenizer | |
from transformers.tokenization_utils_fast import PreTrainedTokenizerFast | |
from configuration_btlm import BTLMConfig | |
from modeling_btlm import BTLMModel, BTLMLMHeadModel | |
AutoConfig.register("btlm", BTLMConfig) | |
AutoModel.register(BTLMConfig, BTLMModel) | |
AutoModelForCausalLM.register(BTLMConfig, BTLMLMHeadModel) | |
AutoTokenizer.register(BTLMConfig, fast_tokenizer_class=PreTrainedTokenizerFast) | |