Register Auto classes

__init__.py

@@ -0,0 +1,30 @@
+# -*- coding: utf-8 -*-
+
+"""
+Hugging Face Model Repository Init File
+
+This file should be renamed to __init__.py and placed at the root of your 
+Hugging Face model repository. It ensures that the model registration code
+is executed when the model is loaded from the Hub.
+"""
+
+# Import the registration code from models/__init__.py
+# This will execute the AutoConfig and AutoModel registration
+try:
+    from models import *
+except ImportError as e:
+    # If there are import issues, try importing the specific registration code
+    from transformers import AutoConfig, AutoModel, AutoModelForCausalLM
+    from models.shared_space_config import SharedSpaceDecoderConfig
+    from models.shared_space_decoder import (
+        SharedSpaceDecoderPreTrainedModel,
+        SharedSpaceDecoderModel,
+    )
+    from layers.task_heads import SharedSpaceDecoderForCausalLM
+    
+    # Register the configuration class with AutoConfig
+    AutoConfig.register("shared_space_decoder", SharedSpaceDecoderConfig)
+    
+    # Register the model classes with AutoModel
+    AutoModel.register(SharedSpaceDecoderConfig, SharedSpaceDecoderModel)
+    AutoModelForCausalLM.register(SharedSpaceDecoderConfig, SharedSpaceDecoderForCausalLM)