refactor(config): 重命名 TransformerConfig 为 ModelConfig

khaosz/config/__init__.py

@@ -1,4 +1,4 @@
-from khaosz.config.model_config import TransformerConfig
+from khaosz.config.model_config import ModelConfig
 from khaosz.config.param_config import BaseModelIO, ModelParameter, Checkpoint, ParameterLoader
 from khaosz.config.schedule_config import ScheduleConfig, CosineScheduleConfig, SGDRScheduleConfig
 from khaosz.config.train_config import TrainConfig
@@ -9,7 +9,7 @@ __all__ = [
     "ModelParameter",
     "Checkpoint",
     "ParameterLoader",
-    "TransformerConfig",
+    "ModelConfig",
     "TrainConfig",
     
     "ScheduleConfig",

khaosz/config/model_config.py

@@ -1,10 +1,10 @@
 import json
 
 from dataclasses import asdict, dataclass
-from typing import  Optional, Self
+from typing import Any, Dict, Optional, Self
 
 @dataclass
-class TransformerConfig:
+class ModelConfig:
     # basic config
     vocab_size: Optional[int] = None
     n_dim: Optional[int] = None
@@ -21,7 +21,7 @@ class TransformerConfig:
     
     def load(self, config_path: str) -> Self:
         with open(config_path, 'r') as f:
-            config: dict = json.load(f)
+            config: Dict[str, Any] = json.load(f)
         for key, value in config.items():
             if hasattr(self, key):
                 setattr(self, key, value)

khaosz/config/param_config.py

@@ -9,7 +9,7 @@ from typing import Any, Dict, List, Optional, Self, Union
 from pathlib import Path
 
 from khaosz.data.tokenizer import BpeTokenizer
-from khaosz.config.model_config import TransformerConfig
+from khaosz.config.model_config import ModelConfig
 from khaosz.model.transformer import Transformer
 
 
@@ -20,11 +20,11 @@ class BaseModelIO:
         self, 
         model: Optional[nn.Module] = None, 
         tokenizer: Optional[BpeTokenizer] = None,
-        config: Optional[TransformerConfig] = None
+        config: Optional[ModelConfig] = None
     ):
         self.model = model
         self.tokenizer = tokenizer or BpeTokenizer()
-        self.config = config or TransformerConfig()
+        self.config = config or ModelConfig()
     
     def _get_file_paths(self, directory: Union[str, Path]) -> dict[str, Path]:
         """Get standardized file paths for model components."""
@@ -79,8 +79,8 @@ class ModelParameter(BaseModelIO):
         default_factory=BpeTokenizer,
         metadata={"help": "Tokenizer for the model."}
     )
-    config: TransformerConfig = field(
-        default_factory=TransformerConfig,
+    config: ModelConfig = field(
+        default_factory=ModelConfig,
         metadata={"help": "Transformer model configuration."}
     )
     
@@ -103,8 +103,8 @@ class Checkpoint(BaseModelIO):
         default_factory=BpeTokenizer,
         metadata={"help": "Tokenizer for the model."}
     )
-    config: TransformerConfig = field(
-        default_factory=TransformerConfig,
+    config: ModelConfig = field(
+        default_factory=ModelConfig,
         metadata={"help": "Transformer model configuration."}
     )
     optimizer_state: Dict[str, Any] = field(
@@ -230,7 +230,7 @@ class ParameterLoader:
     def create_checkpoint(
         model: nn.Module, 
         tokenizer: BpeTokenizer,
-        config: TransformerConfig,
+        config: ModelConfig,
         loss_list: Optional[list[float]] = None,
         optimizer: Optional[optim.Optimizer] = None,
     ) -> Checkpoint: