reafactor: 修改ModelParameter

astrai/config/param_config.py

@@ -1,6 +1,7 @@
 import torch.nn as nn
 import safetensors.torch as st
 
+from contextlib import contextmanager
 from dataclasses import dataclass, field
 from typing import Optional, Self, Union
 from pathlib import Path
@@ -10,12 +11,38 @@ from astrai.config.model_config import ModelConfig
 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: Optional[nn.Module] = field(
-        default=None, metadata={"help": "Transformer model."}
+    model: nn.Module = field(
+        default_factory=nn.Identity, metadata={"help": "Transformer model."}
     )
     tokenizer: BpeTokenizer = field(
         default_factory=BpeTokenizer, metadata={"help": "Tokenizer for the model."}
@@ -41,10 +68,13 @@ class BaseModelIO:
 
         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]) -> Self:
+    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)
 
@@ -52,7 +82,8 @@ class BaseModelIO:
         self.tokenizer.load(str(paths["tokenizer"]))
 
         if self.model is None:
-            self.model = Transformer(self.config)
+            with disable_random_init(enable=disable_init):
+                self.model = Transformer(self.config)
 
         if paths["model"].exists():
             state_dict = st.load_file(str(paths["model"]))
@@ -76,6 +107,8 @@ class ModelParameter(BaseModelIO):
         instance.save_components(save_dir)
 
     @classmethod
-    def load(cls, load_dir: Union[str, Path]) -> "ModelParameter":
+    def load(
+        cls, load_dir: Union[str, Path], disable_init: bool = False
+    ) -> "ModelParameter":
         instance = cls()
-        return instance.load_components(load_dir)
+        return instance.load_components(load_dir, disable_init=disable_init)

astrai/inference/__init__.py

@@ -1,5 +1,4 @@
 from astrai.inference.core import (
-    disable_random_init,
     GeneratorCore,
     EmbeddingEncoderCore,
     KVCacheManager,
@@ -15,7 +14,6 @@ from astrai.inference.generator import (
 )
 
 __all__ = [
-    "disable_random_init",
     "GeneratorCore",
     "EmbeddingEncoderCore",
     "KVCacheManager",

astrai/inference/core.py

@@ -1,8 +1,6 @@
 import torch
-import torch.nn as nn
 
 from torch import Tensor
-from contextlib import contextmanager
 from typing import Any, Callable, List, Tuple, Union, Optional, Self
 from astrai.config import ModelParameter, ModelConfig
 
@@ -55,31 +53,6 @@ def apply_sampling_strategies(
     return logits
 
 
-@contextmanager
-def disable_random_init():
-    init_functions = [
-        "xavier_normal_",
-        "xavier_uniform_",
-        "kaiming_normal_",
-        "kaiming_uniform_",
-        "zeros_",
-        "ones_",
-        "constant_",
-        "normal_",
-        "uniform_",
-    ]
-    original_funcs = {}
-    for name in init_functions:
-        if hasattr(nn.init, name):
-            original_funcs[name] = getattr(nn.init, name)
-            setattr(nn.init, name, lambda *args, **kwargs: None)
-    try:
-        yield
-    finally:
-        for name, orig_func in original_funcs.items():
-            setattr(nn.init, name, orig_func)
-
-
 class GeneratorCore:
     def __init__(self, parameter: ModelParameter):
         self.model = parameter.model