File size: 5,080 Bytes
d91766b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 | from __future__ import annotations
from typing import Any, Callable
from diffulex.config import Config
from diffulex.utils.loader import load_model
_NOT_PROVIDED = object()
RegistryEntry = tuple[Callable[[Any], Any] | type | None, bool]
class AutoModelForDiffusionLM:
"""Factory and registry for diffusion language models."""
MODEL_MAPPING: dict[str, RegistryEntry] = {}
@classmethod
def _ensure_registry_populated(cls) -> None:
"""Best-effort import of built-in models to populate the registry.
This is intentionally defensive for multi-process / spawn execution where
import side-effects can be sensitive to initialization order.
"""
if cls.MODEL_MAPPING:
return
try:
import importlib
import pkgutil
import warnings
import diffulex.model as model_pkg
excluded = {"auto_model", "__init__"}
for mod in pkgutil.iter_modules(model_pkg.__path__):
name = mod.name
if name in excluded or mod.ispkg:
continue
try:
importlib.import_module(f"diffulex.model.{name}")
except Exception as e:
warnings.warn(
f"Failed to import diffulex.model.{name} during registry auto-discovery: {e!r}",
RuntimeWarning,
)
except Exception as e:
# Don't fail hard here; the caller will raise with available models.
try:
import warnings
warnings.warn(
f"Model registry auto-discovery failed: {e!r}",
RuntimeWarning,
)
except Exception:
pass
@classmethod
def register(
cls,
model_name: str,
model_class: Callable[[Any], Any] | type | None = _NOT_PROVIDED,
*,
use_full_config: bool = False,
exist_ok: bool = False,
):
"""Register a model factory or class under ``model_name``.
When ``model_class`` is omitted this method returns a decorator.
Args:
model_name: Key used to retrieve the model.
model_class: Callable or class that builds the model instance.
use_full_config: Pass the entire :class:`Config` to the factory
instead of ``config.hf_config``.
exist_ok: Allow overriding an existing registration.
"""
if not isinstance(model_name, str) or not model_name:
raise ValueError("model_name must be a non-empty string.")
if model_class is _NOT_PROVIDED:
def decorator(model_cls):
cls._register(
model_name,
model_cls,
use_full_config=use_full_config,
exist_ok=exist_ok,
)
return model_cls
return decorator
cls._register(model_name, model_class, use_full_config=use_full_config, exist_ok=exist_ok)
return model_class
@classmethod
def _register(
cls,
model_name: str,
model_class: Callable[[Any], Any] | type | None,
*,
use_full_config: bool,
exist_ok: bool,
) -> None:
if not exist_ok and model_name in cls.MODEL_MAPPING:
raise ValueError(f"Model '{model_name}' is already registered.")
cls.MODEL_MAPPING[model_name] = (model_class, use_full_config)
@classmethod
def unregister(cls, model_name: str) -> None:
cls.MODEL_MAPPING.pop(model_name, None)
@classmethod
def available_models(cls) -> tuple[str, ...]:
return tuple(sorted(cls.MODEL_MAPPING))
@classmethod
def from_config(cls, config: Config):
try:
factory, use_full_config = cls.MODEL_MAPPING[config.model_name]
except KeyError as err:
# Spawn/multi-process execution can hit this before side-effect imports
# have populated the registry. Try a best-effort discovery once.
cls._ensure_registry_populated()
if config.model_name in cls.MODEL_MAPPING:
factory, use_full_config = cls.MODEL_MAPPING[config.model_name]
else:
available = ", ".join(cls.available_models()) or "<none>"
raise ValueError(
f"Model '{config.model_name}' is not registered. Available models: {available}."
) from err
if factory is None:
raise ValueError(f"Model '{config.model_name}' is reserved but not implemented yet.")
init_arg = config if use_full_config else config.hf_config
if init_arg is None:
raise ValueError("Config.hf_config must be initialized before building the model.")
model = factory(init_arg)
return load_model(model, config)
# Backwards compatibility with the old name while callers migrate.
AutoModelLM = AutoModelForDiffusionLM
|