refactor(config): 重命名 TransformerConfig 为 ModelConfig

khaosz/__init__.py

@@ -3,7 +3,7 @@ __author__ = "ViperEkura"
 
 from khaosz.api import Khaosz
 from khaosz.config import (
-    TransformerConfig,
+    ModelConfig,
     ParameterLoader,
     TrainConfig,
 )
@@ -41,7 +41,7 @@ __all__ = [
     "SemanticTextSplitter",
     "PriorityTextSplitter",
     
-    "TransformerConfig",
+    "ModelConfig",
     "ParameterLoader",
     "TrainConfig",
     

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:

khaosz/inference/core.py

@@ -1,7 +1,7 @@
 import torch
 from torch import Tensor 
 from typing import Any, Callable, List, Tuple, Union, Optional, Self
-from khaosz.config import ModelParameter, TransformerConfig
+from khaosz.config import ModelParameter, ModelConfig
 
 
 def apply_sampling_strategies(
@@ -187,7 +187,7 @@ class EmbeddingEncoderCore:
 class KVCacheManager:
     def __init__(
         self, 
-        config: TransformerConfig, 
+        config: ModelConfig, 
         batch_size: int,
         device: torch.device = "cuda", 
         dtype: torch.dtype = torch.bfloat16

khaosz/model/transformer.py

@@ -3,7 +3,7 @@ import torch.nn as nn
 
 from torch import Tensor
 from typing import Any, Mapping, Optional, Tuple
-from khaosz.config.model_config import TransformerConfig
+from khaosz.config.model_config import ModelConfig
 from khaosz.model.module import Embedding, DecoderBlock, Linear, RMSNorm, RotaryEmbedding
 
 
@@ -61,7 +61,7 @@ def process_attention_mask(
 
 
 class Transformer(nn.Module):
-    def __init__(self, config: TransformerConfig):
+    def __init__(self, config: ModelConfig):
         super().__init__()
         self.config = config
         self.rotary_embeding = RotaryEmbedding(config.n_dim // config.n_head, config.m_len)