Skip to content

base

ModelRegistry

An object which is able to return model classes based on given model IDs and model types.

Attributes:

Name Type Description
registry_dict dict

A dictionary mapping model types to model classes.
Source code in inference/core/registries/base.py 
 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
33
34
35
36
37
class ModelRegistry:
    """An object which is able to return model classes based on given model IDs and model types.

    Attributes:
        registry_dict (dict): A dictionary mapping model types to model classes.
    """

    def __init__(self, registry_dict) -> None:
        """Initializes the ModelRegistry with the given dictionary of registered models.

        Args:
            registry_dict (dict): A dictionary mapping model types to model classes.
        """
        self.registry_dict = registry_dict

    def get_model(self, model_type: str, model_id: str) -> Model:
        """Returns the model class based on the given model type.

        Args:
            model_type (str): The type of the model to be retrieved.
            model_id (str): The ID of the model to be retrieved (unused in the current implementation).

        Returns:
            Model: The model class corresponding to the given model type.

        Raises:
            ModelNotRecognisedError: If the model_type is not found in the registry_dict.
        """
        if model_type not in self.registry_dict:
            raise ModelNotRecognisedError(
                f"Could not find model of type: {model_type} in configured registry."
            )
        return self.registry_dict[model_type]

__init__(registry_dict)

Initializes the ModelRegistry with the given dictionary of registered models.

Parameters:

Name Type Description Default
registry_dict dict

A dictionary mapping model types to model classes.

 required
Source code in inference/core/registries/base.py 
12
13
14
15
16
17
18
def __init__(self, registry_dict) -> None:
    """Initializes the ModelRegistry with the given dictionary of registered models.

    Args:
        registry_dict (dict): A dictionary mapping model types to model classes.
    """
    self.registry_dict = registry_dict

get_model(model_type, model_id)

Returns the model class based on the given model type.

Parameters:

Name Type Description Default
model_type str

The type of the model to be retrieved.

 required
model_id str

The ID of the model to be retrieved (unused in the current implementation).

 required

Returns:

Name Type Description
Model Model

The model class corresponding to the given model type.

Raises:

Type Description
ModelNotRecognisedError

If the model_type is not found in the registry_dict.
Source code in inference/core/registries/base.py 
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
def get_model(self, model_type: str, model_id: str) -> Model:
    """Returns the model class based on the given model type.

    Args:
        model_type (str): The type of the model to be retrieved.
        model_id (str): The ID of the model to be retrieved (unused in the current implementation).

    Returns:
        Model: The model class corresponding to the given model type.

    Raises:
        ModelNotRecognisedError: If the model_type is not found in the registry_dict.
    """
    if model_type not in self.registry_dict:
        raise ModelNotRecognisedError(
            f"Could not find model of type: {model_type} in configured registry."
        )
    return self.registry_dict[model_type]