from transformers import AutoConfig, AutoModel from src.models.neuroblast import NeuroBLASTConfig, NeuroBLASTForCausalLM AutoConfig.register("neuroblast", NeuroBLASTForCausalLM) AutoModel.register(NeuroBLASTConfig, NeuroBLASTForCausalLM)