registration.py

from transformers import AutoConfig, AutoModel
from src.models.neuroblast import NeuroBLASTConfig, NeuroBLASTForCausalLM

AutoConfig.register("neuroblast", NeuroBLASTForCausalLM)
AutoModel.register(NeuroBLASTConfig, NeuroBLASTForCausalLM)