feat: 实现模型动态注册机制

astrai/config/__init__.py

@@ -1,11 +1,7 @@
 from astrai.config.model_config import ModelConfig
-from astrai.config.param_config import BaseModelIO, ModelParameter
 from astrai.config.train_config import TrainConfig
 
 __all__ = [
-    # Base I/O
-    "BaseModelIO",
-    "ModelParameter",
     # Model configuration
     "ModelConfig",
     "TrainConfig",

astrai/config/model_config.py

@@ -6,6 +6,7 @@ from typing import Optional, Self
 @dataclass
 class ModelConfig:
     # basic config
+    model_type: Optional[str] = None
     vocab_size: Optional[int] = None
     dim: Optional[int] = None
 

astrai/config/param_config.py

@@ -1,114 +0,0 @@
-from contextlib import contextmanager
-from dataclasses import dataclass, field
-from pathlib import Path
-from typing import Self, Union
-
-import safetensors.torch as st
-import torch.nn as nn
-
-from astrai.config.model_config import ModelConfig
-from astrai.tokenize import BpeTokenizer
-from astrai.model.transformer import Transformer
-
-
-@contextmanager
-def disable_random_init(enable: bool = True):
-    init_functions = [
-        "xavier_normal_",
-        "xavier_uniform_",
-        "kaiming_normal_",
-        "kaiming_uniform_",
-        "zeros_",
-        "ones_",
-        "constant_",
-        "normal_",
-        "uniform_",
-    ]
-    original_funcs = {}
-    for name in init_functions:
-        if enable and hasattr(nn.init, name):
-            original_funcs[name] = getattr(nn.init, name)
-            setattr(nn.init, name, lambda *args, **kwargs: None)
-    try:
-        yield
-    finally:
-        if enable:
-            for name, orig_func in original_funcs.items():
-                setattr(nn.init, name, orig_func)
-
-
-@dataclass
-class BaseModelIO:
-    """Base class for model I/O operations."""
-
-    model: nn.Module = field(
-        default_factory=nn.Identity, metadata={"help": "Transformer model."}
-    )
-    tokenizer: BpeTokenizer = field(
-        default_factory=BpeTokenizer, metadata={"help": "Tokenizer for the model."}
-    )
-    config: ModelConfig = field(
-        default_factory=ModelConfig,
-        metadata={"help": "Transformer model configuration."},
-    )
-
-    def _get_file_paths(self, directory: Union[str, Path]) -> dict[str, Path]:
-        """Get standardized file paths for model components."""
-        dir_path = Path(directory)
-        return {
-            "model": dir_path / "model.safetensors",
-            "config": dir_path / "config.json",
-            "tokenizer": dir_path / "tokenizer.json",
-        }
-
-    def save_components(self, save_dir: Union[str, Path]):
-        """Save core model components."""
-        paths = self._get_file_paths(save_dir)
-        paths["model"].parent.mkdir(parents=True, exist_ok=True)
-
-        if self.model is not None:
-            st.save_file(self.model.state_dict(), str(paths["model"]))
-
-        self.config.save(str(paths["config"]))
-        self.tokenizer.save(str(paths["tokenizer"]))
-
-    def load_components(
-        self, load_dir: Union[str, Path], disable_init: bool = False
-    ) -> Self:
-        """Load core model components."""
-        paths = self._get_file_paths(load_dir)
-
-        self.config.load(str(paths["config"]))
-        self.tokenizer.load(str(paths["tokenizer"]))
-
-        if isinstance(self.model, nn.Identity):
-            with disable_random_init(enable=disable_init):
-                self.model = Transformer(self.config)
-
-        if paths["model"].exists():
-            state_dict = st.load_file(str(paths["model"]))
-            self.model.load_state_dict(state_dict)
-
-        return self
-
-    def to(self, *args, **kwargs) -> "BaseModelIO":
-        """Move model to device."""
-        if self.model is not None:
-            self.model.to(*args, **kwargs)
-        return self
-
-
-@dataclass
-class ModelParameter(BaseModelIO):
-    """Container for model parameters with serialization capabilities."""
-
-    @classmethod
-    def save(cls, instance: "ModelParameter", save_dir: Union[str, Path]):
-        instance.save_components(save_dir)
-
-    @classmethod
-    def load(
-        cls, load_dir: Union[str, Path], disable_init: bool = False
-    ) -> "ModelParameter":
-        instance = cls()
-        return instance.load_components(load_dir, disable_init=disable_init)