File size: 864 Bytes
69928d6 | 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 | """Model registry — maps config name to model class."""
from __future__ import annotations
from typing import Any
import torch.nn as nn
MODEL_REGISTRY: dict[str, type[nn.Module]] = {}
def register_model(name: str):
"""Decorator to register a model class."""
def decorator(cls):
MODEL_REGISTRY[name] = cls
return cls
return decorator
def get_model(name: str, **kwargs: Any) -> nn.Module:
"""Instantiate a registered model by name."""
if name not in MODEL_REGISTRY:
raise ValueError(f"Unknown model '{name}'. Available: {list(MODEL_REGISTRY.keys())}")
return MODEL_REGISTRY[name](**kwargs)
# Auto-import submodules to trigger registration
from softasp_solver.sudoku import baselines as _sudoku_baselines # noqa: F401, E402
from softasp_solver.sudoku import softasp_sudoku as _sudoku # noqa: F401, E402
|