ViperEkura/AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

refactor: simplify BaseFactory and separate ModelFactory from AutoModel

Browse Source 
- 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:
ViperEkura
2026-07-30 09:38:20 +08:00
parent 22cf798d81
commit fc47319240
8 changed files with 59 additions and 65 deletions
Download Patch File Download Diff File Expand all files Collapse all files
astrai/factory.py
+41 -35
View File
@@ -13,41 +13,63 @@ from typing import (
Type,
TypeVar,
Union,
get_args,
get_origin,
)
from typing import get_args as _get_args
from typing import get_origin as _get_origin
T = TypeVar("T")
def _resolve_type(
def _resolve_base_type(
arg: Union[Type, str, ForwardRef], factory_cls: type
) -> Optional[Type]:
"""Resolve a generic type-arg (str forward-ref, ForwardRef, or class)."""
if not isinstance(arg, (str, ForwardRef)):
"""Resolve the generic type-arg T to a concrete class.
- Concrete class (``BaseFactory[MyBase]``): returned directly.
- Forward reference (``BaseFactory["MyBase"]``): ``Base["X"]``
produces a ``ForwardRef("X")`` at class-creation time. We
extract the name and evaluate it in the factory module's
global namespace — the same mechanism ``typing.get_type_hints``
uses internally.
"""
if isinstance(arg, type):
return arg
name = arg if isinstance(arg, str) else arg.__forward_arg__
if name == factory_cls.__name__:
return factory_cls
if isinstance(arg, str):
name = arg
elif isinstance(arg, ForwardRef):
name = arg.__forward_arg__
else:
return None
mod = sys.modules.get(factory_cls.__module__)
if mod is None:
return None
ns = vars(mod)
try:
return eval(name, vars(mod)) # noqa: S307
except NameError:
return None
if isinstance(arg, ForwardRef):
return arg._evaluate(ns, None, recursive_guard=frozenset())
return ns.get(name)
def _validate_component(component_cls: Type, base: Optional[Type]) -> None:
"""Validate that *component_cls* inherits from *base*.
No-op when *base* is ``None`` (e.g. forward-ref resolution failed).
"""
if base is not None and not issubclass(component_cls, base):
raise TypeError(f"{component_cls.__name__} must inherit from {base.__name__}")
class BaseFactory(ABC, Generic[T]):
"""Generic factory with decorator-based component registration.
"""Generic factory with decorator-based registration.
Create a factory by subclassing with the desired base type::
class MyFactory(BaseFactory[MyBase]):
pass
Register components with the ``register`` decorator::
@MyFactory.register("custom")
class CustomComponent(MyBase):
...
@@ -64,13 +86,10 @@ class BaseFactory(ABC, Generic[T]):
def __init_subclass__(cls, **kwargs):
super().__init_subclass__(**kwargs)
for orig_base in getattr(cls, "__orig_bases__", ()):
if _get_origin(orig_base) is BaseFactory:
(arg,) = _get_args(orig_base)
if get_origin(orig_base) is BaseFactory:
(arg,) = get_args(orig_base)
cls._entries = {}
try:
cls._component_base = _resolve_type(arg, cls)
except Exception:
cls._component_base = None
cls._component_base = _resolve_base_type(arg, cls)
return
@classmethod
@@ -82,7 +101,7 @@ class BaseFactory(ABC, Generic[T]):
"""
def decorator(component_cls: Type[T]) -> Type[T]:
cls._validate_component(component_cls)
_validate_component(component_cls, cls._component_base)
if name in cls._entries:
raise ValueError(f"Component '{name}' is already registered")
cls._entries[name] = component_cls
@@ -95,12 +114,11 @@ class BaseFactory(ABC, Generic[T]):
"""Create a component instance by name, filtering kwargs to match
the component's ``__init__`` signature.
"""
entry = cls._entries.get(name)
if entry is None:
component_cls = cls._entries.get(name)
if component_cls is None:
raise ValueError(
f"Unknown component: '{name}'. Supported types: {sorted(cls._entries)}"
)
component_cls = entry
sig = inspect.signature(component_cls.__init__)
has_var_kwargs = any(
p.kind == inspect.Parameter.VAR_KEYWORD for p in sig.parameters.values()
@@ -114,18 +132,6 @@ class BaseFactory(ABC, Generic[T]):
kwargs = {k: v for k, v in kwargs.items() if k in valid}
return component_cls(*args, **kwargs)
@classmethod
def _validate_component(cls, component_cls: Type[T]):
"""Validate the decorated class inherits from the factory's base type.
Override for custom validation beyond ``issubclass``.
"""
base = cls._component_base
if base is not None and not issubclass(component_cls, base):
raise TypeError(
f"{component_cls.__name__} must inherit from {base.__name__}"
)
@classmethod
def get_component_class(cls, name: str) -> Type[T]:
"""Get the registered component class without instantiating it."""
astrai/model/automodel.py
+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)
astrai/model/encoder.py
+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)
astrai/model/transformer.py
+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."""
astrai/tokenize/tokenizer.py
-13
View File
@@ -20,8 +20,6 @@ Messages = List[Message]
class AutoTokenizer:
"""Base tokenizer class with automatic loading support"""
TOKENIZER_CLASSES = {} # Registry for auto-loading
def __init__(
self,
path: Optional[Union[str, Path]] = None,
@@ -108,17 +106,6 @@ class AutoTokenizer:
with open(save_path / "tokenizer_config.json", "w", encoding="utf-8") as f:
json.dump(config, f, ensure_ascii=False, indent=2)
@classmethod
def register_tokenizer(cls, name: str, tokenizer_class: type):
"""
Register a new tokenizer class.
Args:
name: Name to register the tokenizer class under
tokenizer_class: The tokenizer class to register
"""
cls.TOKENIZER_CLASSES[name] = tokenizer_class
def encode(
self,
tokens: Union[str, List[str]],
astrai/trainer/strategy.py
+2 -2
View File
@@ -501,5 +501,5 @@ class GRPOStrategy(BaseStrategy):
# Factory aliases: online variants use the same strategy class; the
# ``RolloutRunner`` is injected by ``TrainContextBuilder`` to enable
# online mode, so no separate subclass is needed.
StrategyFactory._entries["online_grpo"] = GRPOStrategy
StrategyFactory._entries["online_dpo"] = DPOStrategy
StrategyFactory.register("online_grpo")(GRPOStrategy)
StrategyFactory.register("online_dpo")(DPOStrategy)
tests/module/test_encoder.py
+3 -3
View File
@@ -7,7 +7,7 @@ import safetensors.torch as st
import torch
from astrai.config.model_config import EncoderConfig
from astrai.model.automodel import AutoModel
from astrai.model.automodel import ModelFactory
from astrai.model.encoder import EmbeddingEncoder
from tests.helpers import TINY_CONFIG, assert_state_dicts_equal
@@ -69,8 +69,8 @@ def test_encoder_normalize(device):
def test_encoder_register():
assert AutoModel.is_registered("embedding")
cls = AutoModel.get_component_class("embedding")
assert ModelFactory.is_registered("embedding")
cls = ModelFactory.get_component_class("embedding")
assert cls is EmbeddingEncoder
tests/trainer/test_online_strategy.py
+2 -2
View File
@@ -92,8 +92,8 @@ def _make_dpo(device, executor=None):
def test_factory_registers_online_aliases():
assert StrategyFactory.is_registered("online_grpo")
assert StrategyFactory.is_registered("online_dpo")
assert StrategyFactory._entries["online_grpo"] is GRPOStrategy
assert StrategyFactory._entries["online_dpo"] is DPOStrategy
assert StrategyFactory.get_component_class("online_grpo") is GRPOStrategy
assert StrategyFactory.get_component_class("online_dpo") is DPOStrategy
def test_grpo_supports_online(device):