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
| 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]
|