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

astrai/model/automodel.py

@@ -0,0 +1,134 @@
+"""
+AutoModel base class for model loading and saving.
+"""
+
+import torch.nn as nn
+import safetensors.torch as st
+
+from pathlib import Path
+from contextlib import contextmanager
+from typing import Self, Union, Dict, Type
+
+from astrai.config import ModelConfig
+
+
+@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)
+
+
+class AutoModel(nn.Module):
+    """
+    Autoregressive language model base class.
+    Provides model loading/saving and generation capabilities.
+    """
+
+    # Model registry - stored as class attribute
+    _registry: Dict[str, Type["AutoModel"]] = {}
+
+    def __init__(self, config: ModelConfig):
+        super().__init__()
+        self.config = config
+
+    @classmethod
+    def register(cls, model_type: str):
+        """
+        Class method decorator to register model type.
+
+        Usage:
+            @AutoModel.register('transformer')
+            class Transformer(AutoModel):
+                ...
+        """
+
+        def decorator(sub_cls: Type["AutoModel"]) -> Type["AutoModel"]:
+            cls._registry[model_type.lower()] = sub_cls
+            return sub_cls
+
+        return decorator
+
+    @classmethod
+    def get_model_class(cls, model_type: str) -> Type["AutoModel"]:
+        """Get model class by model_type string."""
+        model_type = model_type.lower()
+        if model_type not in cls._registry:
+            available = list(cls._registry.keys())
+            raise ValueError(
+                f"Unknown model_type: {model_type}. Available: {available}"
+            )
+        return cls._registry[model_type]
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        path: Union[str, Path],
+        disable_random_init: bool = True,
+    ) -> nn.Module:
+
+        model_path = Path(path)
+
+        # Load config
+        config = ModelConfig()
+        config_path = model_path / "config.json"
+        if config_path.exists():
+            config.load(str(config_path))
+        else:
+            raise FileNotFoundError(f"Config file not found: {config_path}")
+
+        # If called from base class, use model_type to determine actual model class
+        if cls is AutoModel:
+            model_type = config.model_type or "transformer"
+            actual_cls = cls.get_model_class(model_type)
+        else:
+            raise ValueError(
+                f"Cannot call from_pretrained() on subclass {cls.__name__}"
+            )
+
+        with _disable_random_init(enable=disable_random_init):
+            model = actual_cls(config)
+
+        # Load weights
+        weights_path = model_path / "model.safetensors"
+        if weights_path.exists():
+            state_dict = st.load_file(str(weights_path))
+            model.load_state_dict(state_dict, strict=False)
+
+        return model
+
+    def save_pretrained(
+        self,
+        save_directory: Union[str, Path],
+    ) -> None:
+        save_path = Path(save_directory)
+        save_path.mkdir(parents=True, exist_ok=True)
+
+        # Save config
+        self.config.save(str(save_path / "config.json"))
+
+        # Save weights
+        st.save_file(self.state_dict(), str(save_path / "model.safetensors"))
+
+    def to(self, *args, **kwargs) -> Self:
+        """Move model to device/dtype."""
+        return super().to(*args, **kwargs)