registry

registry

registry.py Decorator-based registry for jobs, datasets, and evaluations.

Decorators (@job, @dataset, @evaluation) can be imported cheaply and used to register classes at definition time — no heavy submodule imports happen until the registry dicts are actually read.

Usage

from theseus.registry import job, dataset, evaluation

@job("gpt/train/pretrain") class PretrainGPT(BaseTrainer[...]): ...

@dataset("alpaca") class Alpaca(ChatTemplateDataset): ...

@evaluation("bbq") class BBQEval(RolloutEvaluation): ...

User-defined jobs in scripts are recognized automatically — decorate with @job/@dataset/@evaluation and the class will coexist with built-in entries the moment any code reads from the registry.

ensure_registered() -> None

Trigger registration of all built-in decorated classes.

Safe to call multiple times — only the first call imports the submodules. Any classes already registered via decorators (e.g. user-defined jobs) are preserved.

job(key: str) -> Callable[[T], T]

Register a job class under the given key.

dataset(key: str) -> Callable[[T], T]

Register a dataset class under the given key.

evaluation(key: str) -> Callable[[T], T]

Register an evaluation callable under the given key.

PerplexityEvaluation subclasses must use a key ending in _ppl so downstream consumers (e.g. the boundary-eval bar plots) can group unbounded perplexity metrics away from 0-1 scored ones. PerplexityComparisonEvaluation returns an accuracy, not a perplexity, and is intentionally exempt.