Skip to content

Model factory

models

Model factory.

Importing this package registers every model file by name (each model module ends with register(...)); get_model(name) returns the class and list_models() lists the registered names. Adding a model is one new file + one register call listed below — no other edits.

The model modules are imported here for their registration side effect. The import block grows by one line per model.

base_model_name

base_model_name(name)

Return the canonical registered base model name from a <base>@<tag> variant string.

Source code in src/scale_forecasting/models/__init__.py
def base_model_name(name: str) -> str:
    """Return the canonical registered base model name from a ``<base>@<tag>`` variant string."""
    base = name.split("@", 1)[0].strip()
    return _MODEL_ALIASES.get(base, base)

filter_available_models

filter_available_models(models)

Filter models to those whose upstream package is installed in this environment.

Logs a warning for each omitted model so restricted environments can run a broad model list without failing on uninstalled optional packages.

Source code in src/scale_forecasting/models/__init__.py
def filter_available_models(models: Sequence[str]) -> list[str]:
    """Filter ``models`` to those whose upstream package is installed in this environment.

    Logs a warning for each omitted model so restricted environments can run a broad
    model list without failing on uninstalled optional packages.
    """
    kept: list[str] = []
    omitted: list[str] = []
    for name in models:
        cls = get_model(name)
        if cls.is_available():
            kept.append(name)
        else:
            omitted.append(f"{name} (requires {cls.package})")
    if omitted:
        _log.warning(
            "Omitting %d model(s) whose optional packages are not installed: %s",
            len(omitted),
            ", ".join(omitted),
        )
    return kept

get_model

get_model(name)

Return the registered model class for name (or <base>@<tag> variant).

Raises ModelError with the available names when name is unknown.

Source code in src/scale_forecasting/models/__init__.py
def get_model(name: str) -> type[BaseModel]:
    """Return the registered model class for ``name`` (or ``<base>@<tag>`` variant).

    Raises ``ModelError`` with the available names when ``name`` is unknown.
    """
    base = base_model_name(name)
    try:
        return _REGISTRY[base]
    except KeyError:
        known = ", ".join(sorted(_REGISTRY)) or "(none registered)"
        raise ModelError(f"unknown model '{name}'; registered models: {known}") from None

list_models

list_models(*, available_only=False)

All registered model names, sorted.

Pass available_only=True to restrict the returned list to models whose upstream Python package is installed in the current environment.

Source code in src/scale_forecasting/models/__init__.py
def list_models(*, available_only: bool = False) -> list[str]:
    """All registered model names, sorted.

    Pass ``available_only=True`` to restrict the returned list to models whose
    upstream Python package is installed in the current environment.
    """
    if not available_only:
        return sorted(_REGISTRY)
    return sorted(name for name, cls in _REGISTRY.items() if cls.is_available())