diff --git a/astrai/factory.py b/astrai/factory.py index 9d9e5b7..c1cbde4 100644 --- a/astrai/factory.py +++ b/astrai/factory.py @@ -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.""" diff --git a/astrai/model/automodel.py b/astrai/model/automodel.py index 5459314..6d79a3e 100644 --- a/astrai/model/automodel.py +++ b/astrai/model/automodel.py @@ -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) diff --git a/astrai/model/encoder.py b/astrai/model/encoder.py index 3d2bd4a..1229c9d 100644 --- a/astrai/model/encoder.py +++ b/astrai/model/encoder.py @@ -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) diff --git a/astrai/model/transformer.py b/astrai/model/transformer.py index c194d52..902d8e6 100644 --- a/astrai/model/transformer.py +++ b/astrai/model/transformer.py @@ -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.""" diff --git a/astrai/tokenize/tokenizer.py b/astrai/tokenize/tokenizer.py index 3201dda..5fe1357 100644 --- a/astrai/tokenize/tokenizer.py +++ b/astrai/tokenize/tokenizer.py @@ -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]], diff --git a/astrai/trainer/strategy.py b/astrai/trainer/strategy.py index 9ddc422..6f7bb2f 100644 --- a/astrai/trainer/strategy.py +++ b/astrai/trainer/strategy.py @@ -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) diff --git a/tests/module/test_encoder.py b/tests/module/test_encoder.py index 3b3b0a9..df86071 100644 --- a/tests/module/test_encoder.py +++ b/tests/module/test_encoder.py @@ -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 diff --git a/tests/trainer/test_online_strategy.py b/tests/trainer/test_online_strategy.py index 91544ba..27e9b94 100644 --- a/tests/trainer/test_online_strategy.py +++ b/tests/trainer/test_online_strategy.py @@ -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):