refactor: simplify BaseFactory and separate ModelFactory from AutoModel

- Extract _resolve_base_type and _validate_component as module-level helpers
- Replace ForwardRef._evaluate private API with eval in module namespace
- Remove broad except Exception in __init_subclass__, _component_base always set
- Replace direct _entries mutation in strategy.py with register() call form
- Remove dead TOKENIZER_CLASSES registry from AutoTokenizer
- Extract ModelFactory(BaseFactory[nn.Module]) as pure factory
- AutoModel now inherits only nn.Module, no factory state
- Move @AutoModel.register to @ModelFactory.register in transformer.py and encoder.py
This commit is contained in:
2026-07-30 09:38:20 +08:00
parent 22cf798d81
commit fc47319240
8 changed files with 59 additions and 65 deletions
+7 -6
View File
@@ -40,11 +40,12 @@ def _disable_random_init(enable: bool = True):
setattr(nn.init, n, fn)
class AutoModel(BaseFactory["AutoModel"], nn.Module):
"""
Autoregressive language model base class.
Provides model loading/saving, registration, and generation.
"""
class ModelFactory(BaseFactory[nn.Module]):
"""Pure factory for model dispatch, separated from nn.Module state."""
class AutoModel(nn.Module):
"""Model base class with loading/saving and generation."""
def __init__(self, config: BaseModelConfig):
super().__init__()
@@ -68,7 +69,7 @@ class AutoModel(BaseFactory["AutoModel"], nn.Module):
config = ConfigFactory.load(raw)
model_type = config.model_type or "autoregressive_lm"
actual_cls = AutoModel.get_component_class(model_type)
actual_cls = ModelFactory.get_component_class(model_type)
with _disable_random_init(enable=disable_random_init):
model = actual_cls(config)
+2 -2
View File
@@ -5,7 +5,7 @@ import torch.nn as nn
from torch import Tensor
from astrai.config.model_config import EncoderConfig
from astrai.model.automodel import AutoModel
from astrai.model.automodel import AutoModel, ModelFactory
from astrai.model.components.decoder_block import DecoderBlock
from astrai.model.components.embedding import Embedding
from astrai.model.components.norm import RMSNorm
@@ -13,7 +13,7 @@ from astrai.model.components.rope import RotaryEmbedding
from astrai.model.transformer import process_attention_mask
@AutoModel.register("embedding")
@ModelFactory.register("embedding")
class EmbeddingEncoder(AutoModel):
def __init__(self, config: EncoderConfig):
super().__init__(config)
+2 -2
View File
@@ -6,7 +6,7 @@ from torch import Tensor
from astrai.config.model_config import AutoRegressiveLMConfig
from astrai.inference.core.cache import CacheView
from astrai.model.automodel import AutoModel
from astrai.model.automodel import AutoModel, ModelFactory
from astrai.model.components.decoder_block import DecoderBlock
from astrai.model.components.embedding import Embedding
from astrai.model.components.linear import Linear
@@ -26,7 +26,7 @@ def process_attention_mask(
return input_mask
@AutoModel.register("autoregressive_lm")
@ModelFactory.register("autoregressive_lm")
class AutoRegressiveLM(AutoModel):
"""Autoregressive language model with paged KV cache."""