registry
registry
¶
registry.py Decorator-based registry for jobs, analyses, datasets, and evaluations.
Decorators (@job, @analysis, @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, analysis, dataset, evaluation
@job("gpt/train/pretrain") class PretrainGPT(BaseTrainer[...]): ...
@analysis("gpt/analyze/attention") class GPTAttention(AttentionHeatmapAnalysis, GPTPretrain): ...
@dataset("alpaca") class Alpaca(ChatTemplateDataset): ...
@evaluation("bbq") class BBQEval(RolloutEvaluation): ...
User-defined jobs in scripts are recognized automatically — decorate with @job/@analysis/@dataset/@evaluation and the class will coexist with built-in entries the moment any code reads from the registry. Analyses also appear in JOBS; the analysis templates themselves are unregistered.
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.
is_builtin_job(cls: Any) -> bool
¶
Whether a class is defined in the installed Theseus package.
job(key: str) -> Callable[[T], T]
¶
Register a job class under the given key.
analysis(key: str) -> Callable[[T], T]
¶
Register an analysis as both a job and a discoverable analysis.
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. Bits-per-byte
evaluations use _bpb for the same reason.
PerplexityComparisonEvaluation returns an accuracy, not a
perplexity, and is intentionally exempt.