delete config_cls and model_cls, only maintain auto_map

Author: RkGrit

iotdb-core/ainode/iotdb/ainode/core/model/model_info.py

@@ -28,8 +28,6 @@ def __init__(
         category: ModelCategory,
         state: ModelStates,
         model_type: str = "",
-        config_cls: str = "",
-        model_cls: str = "",
         pipeline_cls: str = "",
         repo_id: str = "",
         auto_map: Optional[Dict] = None,
@@ -39,8 +37,6 @@ def __init__(
         self.model_type = model_type
         self.category = category
         self.state = state
-        self.config_cls = config_cls
-        self.model_cls = model_cls
         self.pipeline_cls = pipeline_cls
         self.repo_id = repo_id
         self.auto_map = auto_map  # If exists, indicates it's a Transformers model
@@ -114,19 +110,23 @@ def __repr__(self):
         category=ModelCategory.BUILTIN,
         state=ModelStates.INACTIVE,
         model_type="timer",
-        config_cls="configuration_timer.TimerConfig",
-        model_cls="modeling_timer.TimerForPrediction",
         pipeline_cls="pipeline_timer.TimerPipeline",
         repo_id="thuml/timer-base-84m",
+        auto_map={
+            "AutoConfig": "configuration_timer.TimerConfig",
+            "AutoModelForCausalLM": "modeling_timer.TimerForPrediction",
+        },
     ),
     "sundial": ModelInfo(
         model_id="sundial",
         category=ModelCategory.BUILTIN,
         state=ModelStates.INACTIVE,
         model_type="sundial",
-        config_cls="configuration_sundial.SundialConfig",
-        model_cls="modeling_sundial.SundialForPrediction",
         pipeline_cls="pipeline_sundial.SundialPipeline",
         repo_id="thuml/sundial-base-128m",
+        auto_map={
+            "AutoConfig": "configuration_sundial.SundialConfig",
+            "AutoModelForCausalLM": "modeling_sundial.SundialForPrediction",
+        },
     ),
 }

iotdb-core/ainode/iotdb/ainode/core/model/model_loader.py

@@ -69,22 +69,25 @@ def load_model_from_transformers(model_info: ModelInfo, **model_kwargs):
         model_info.model_id,
     )
 
+    config_str = model_info.auto_map.get("AutoConfig", "")
+    model_str = model_info.auto_map.get("AutoModelForCausalLM", "")
+
     if model_info.category == ModelCategory.BUILTIN:
         module_name = (
             AINodeDescriptor().get_config().get_ain_models_builtin_dir()
             + "."
             + model_info.model_id
         )
-        config_cls = import_class_from_path(module_name, model_info.config_cls)
-        model_cls = import_class_from_path(module_name, model_info.model_cls)
-    elif model_info.model_cls and model_info.config_cls:
+        config_cls = import_class_from_path(module_name, config_str)
+        model_cls = import_class_from_path(module_name, model_str)
+    elif model_str and config_str:
         module_parent = str(Path(model_path).parent.absolute())
         with temporary_sys_path(module_parent):
             config_cls = import_class_from_path(
-                model_info.model_id, model_info.config_cls
+                model_info.model_id, config_str
             )
             model_cls = import_class_from_path(
-                model_info.model_id, model_info.model_cls
+                model_info.model_id, model_str
             )
     else:
         config_cls = AutoConfig.from_pretrained(model_path)