Update __init__.py

__init__.py

@@ -1,8 +1,33 @@
-from .configuration_arlow_gpt import ArlowGPTConfig
-from .modeling_arlow_gpt import ArlowGPT
-# Register the model
-from transformers import AutoConfig, AutoModelForCausalLM
-
-# Register auto classes
-AutoConfig.register("arlow_gpt", ArlowGPTConfig)
-AutoModelForCausalLM.register(ArlowGPTConfig, ArlowGPT)
+from typing import TYPE_CHECKING
+
+from transformers.utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available
+
+_import_structure = {
+    "configuration_arlow_gpt": ["ARLOW_GPT_PRETRAINED_CONFIG_ARCHIVE_MAP", "ArlowGPTConfig"],
+}
+
+try:
+    if not is_torch_available():
+        raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+    pass
+else:
+    _import_structure["modeling_arlow_gpt"] = [
+        "ArlowGPT",
+        "ArlowGPTPreTrainedModel",
+    ]
+
+if TYPE_CHECKING:
+    from .configuration_arlow_gpt import ARLOW_GPT_PRETRAINED_CONFIG_ARCHIVE_MAP, ArlowGPTConfig
+
+    try:
+        if not is_torch_available():
+            raise OptionalDependencyNotAvailable()
+    except OptionalDependencyNotAvailable:
+        pass
+    else:
+        from .modeling_arlow_gpt import ArlowGPT, ArlowGPTPreTrainedModel
+
+else:
+    import sys
+    sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)