File size: 3,490 Bytes
31dc8dc | 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 | from __future__ import annotations
from typing import Any, Callable
from diffulex.config import Config
_NOT_PROVIDED = object()
RegistryEntry = tuple[Callable[[Any], Any] | type | None, bool]
class AutoSampler:
"""Factory and registry for diffusion language model samplers."""
SAMPLER_MAPPING: dict[str, RegistryEntry] = {}
@classmethod
def register(
cls,
sampler_name: str,
sampler_class: Callable[[Any], Any] | type | None = _NOT_PROVIDED,
*,
use_full_config: bool = False,
exist_ok: bool = False,
):
"""Register a sampler factory or class under ``sampler_name``.
When ``sampler_class`` is omitted this method returns a decorator.
Args:
sampler_name: Key used to retrieve the sampler.
sampler_class: Callable or class that builds the sampler 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(sampler_name, str) or not sampler_name:
raise ValueError("sampler_name must be a non-empty string.")
if sampler_class is _NOT_PROVIDED:
def decorator(sampler_cls):
cls._register(
sampler_name,
sampler_cls,
use_full_config=use_full_config,
exist_ok=exist_ok,
)
return sampler_cls
return decorator
cls._register(
sampler_name,
sampler_class,
use_full_config=use_full_config,
exist_ok=exist_ok,
)
return sampler_class
@classmethod
def _register(
cls,
sampler_name: str,
sampler_class: Callable[[Any], Any] | type | None,
*,
use_full_config: bool,
exist_ok: bool,
) -> None:
if not exist_ok and sampler_name in cls.SAMPLER_MAPPING:
raise ValueError(f"Sampler '{sampler_name}' is already registered.")
cls.SAMPLER_MAPPING[sampler_name] = (sampler_class, use_full_config)
@classmethod
def unregister(cls, sampler_name: str) -> None:
cls.SAMPLER_MAPPING.pop(sampler_name, None)
@classmethod
def available_samplers(cls) -> tuple[str, ...]:
return tuple(sorted(cls.SAMPLER_MAPPING))
@classmethod
def from_config(cls, config: Config):
try:
factory, use_full_config = cls.SAMPLER_MAPPING[config.model_name]
except KeyError as err:
available = ", ".join(cls.available_samplers()) or "<none>"
raise ValueError(
f"Sampler '{config.model_name}' is not registered. Available samplers: {available}."
) from err
if factory is None:
raise ValueError(f"Sampler '{config.model_name}' is reserved but not implemented yet.")
sampler = factory(config) if use_full_config else factory()
if sampler is None:
raise ValueError(
"Unsupported sampler configuration for "
f"model_name='{config.model_name}', "
f"decoding_strategy='{config.decoding_strategy}', "
f"sampling_mode='{config.sampling_mode}'."
)
setattr(sampler, "tokenizer_vocab_size", getattr(config, "tokenizer_vocab_size", None))
return sampler
|