Update modeling_neollm.py

modeling_neollm.py

@@ -30,6 +30,7 @@ from transformers.utils.import_utils import (
 )
 from configuration_neollm import NeoLLMConfig
 
+from transformers import AutoConfig, AutoModel, AutoModelForCausalLM
 
 if is_causal_conv1d_available():
     from causal_conv1d import causal_conv1d_fn, causal_conv1d_update
@@ -1016,6 +1017,12 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
             attentions=outputs.attentions,
         )
 
+# ==================== AUTOMODEL REGISTRATION ====================
+
+# Register the configuration and model for AutoClass support
+AutoConfig.register("neollm", NeoLLMConfig)
+AutoModel.register(NeoLLMConfig, NeoLLMModel)
+AutoModelForCausalLM.register(NeoLLMConfig, NeoLLMForCausalLM)
 
 __all__ = [
     "NeoLLMForCausalLM",
@@ -1025,10 +1032,3 @@ __all__ = [
     "FANLayer",
 ]
 
-# ==================== AUTOMODEL REGISTRATION ====================
-from transformers import AutoConfig, AutoModel, AutoModelForCausalLM
-
-# Register the configuration and model for AutoClass support
-AutoConfig.register("neollm", NeoLLMConfig)
-AutoModel.register(NeoLLMConfig, NeoLLMModel)
-AutoModelForCausalLM.register(NeoLLMConfig, NeoLLMForCausalLM)