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:
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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."""
|
||||
|
||||
|
||||
Reference in New Issue
Block a user