oumi.core.registry

oumi.core.registry#

A framework used for registering and accessing objects across Oumi.

class oumi.core.registry.Registry[source]#

Bases: object

__getitem__(args: tuple[str, RegistryType]) → Callable[source]#

Gets a record by name and type.

__repr__() → str[source]#

Defines how this class is properly printed.

clear() → None[source]#

Clears the registry.

contains(name: str, type: RegistryType) → bool[source]#

Indicates whether a record exists in the registry.

get(name: str, type: RegistryType) → Callable | None[source]#

Gets a record by name and type.

get_all(type: RegistryType) → dict[source]#

Gets all records of a specific type.

get_dataset(name: str, subset: str | None = None) → Callable | None[source]#

Gets a record that corresponds to a registered dataset.

get_judge_config(name: str) → Callable | None[source]#

Gets a record that corresponds to a registered judge config.

get_metrics_function(name: str) → Callable | None[source]#

Gets a record that corresponds to a registered metrics function.

get_model(name: str) → Callable | None[source]#

Gets a record that corresponds to a registered model.

get_model_config(name: str) → Callable | None[source]#

Gets a record that corresponds to a registered config.

register(name: str, type: RegistryType, value: Any) → None[source]#

Registers a new record.

class oumi.core.registry.RegistryType(value, names=None, *, module=None, qualname=None, type=None, start=1, boundary=None)[source]#

Bases: Enum

CLOUD = 1#
DATASET = 2#
JUDGE_CONFIG = 6#
METRICS_FUNCTION = 3#
MODEL = 5#
MODEL_CONFIG = 4#
oumi.core.registry.register(registry_name: str, registry_type: RegistryType) → Callable[source]#

Returns function to register decorated obj in the Oumi global registry.

Parameters:

  • registry_name – The name that the object should be registered with.

  • registry_type – The type of object we are registering.

Returns:

Decorator function to register the target object.

oumi.core.registry.register_cloud_builder(registry_name: str) → Callable[source]#

Returns a function to register decorated builder in the Oumi global registry.

Use this decorator to register cloud builder functions in the global registry. A cloud builder function is a function that accepts no arguments and returns an instance of a class that implements the BaseCloud interface.

Parameters:

registry_name – The name that the builder should be registered with.

Returns:

Decorator function to register the target builder.

oumi.core.registry.register_dataset(registry_name: str, subset: str | None = None) → Callable[source]#

Returns function to register decorated obj in the Oumi global registry.

Parameters:

  • registry_name – The name that the object should be registered with.

  • subset – The type of object we are registering.

Returns:

Decorator function to register the target object.

oumi.core.registry.register_judge(registry_name: str) → Callable[source]#

Returns a function to register a judge configuration in the Oumi global registry.

This decorator is used to register judge configuration in the global registry. A judge configuration function typically returns a JudgeConfig object that defines the parameters and attributes for a specific judge.

Parameters:

registry_name – The name under which the judge configuration should be registered.

Returns:

A decorator function that registers the target judge configuration.

Return type:

Callable

Example

@register_judge("my_custom_judge")
def my_judge_config() -> JudgeConfig:
    return JudgeConfig(...)