bikes.io.registries

Savers, loaders, and registers for model registries.

  1"""Savers, loaders, and registers for model registries."""
  2
  3# %% IMPORTS
  4
  5import abc
  6import typing as T
  7
  8import mlflow
  9import pydantic as pdt
 10from mlflow.pyfunc import PyFuncModel, PythonModel, PythonModelContext
 11
 12from bikes.core import models, schemas
 13from bikes.utils import signers
 14
 15# %% TYPES
 16
 17# Results of model registry operations
 18Info: T.TypeAlias = mlflow.models.model.ModelInfo
 19Alias: T.TypeAlias = mlflow.entities.model_registry.ModelVersion
 20Version: T.TypeAlias = mlflow.entities.model_registry.ModelVersion
 21
 22# %% HELPERS
 23
 24
 25def uri_for_model_alias(name: str, alias: str) -> str:
 26    """Create a model URI from a model name and an alias.
 27
 28    Args:
 29        name (str): name of the mlflow registered model.
 30        alias (str): alias of the registered model.
 31
 32    Returns:
 33        str: model URI as "models:/name@alias".
 34    """
 35    return f"models:/{name}@{alias}"
 36
 37
 38def uri_for_model_version(name: str, version: int) -> str:
 39    """Create a model URI from a model name and a version.
 40
 41    Args:
 42        name (str): name of the mlflow registered model.
 43        version (int): version of the registered model.
 44
 45    Returns:
 46        str: model URI as "models:/name/version."
 47    """
 48    return f"models:/{name}/{version}"
 49
 50
 51def uri_for_model_alias_or_version(name: str, alias_or_version: str | int) -> str:
 52    """Create a model URi from a model name and an alias or version.
 53
 54    Args:
 55        name (str): name of the mlflow registered model.
 56        alias_or_version (str | int): alias or version of the registered model.
 57
 58    Returns:
 59        str: model URI as "models:/name@alias" or "models:/name/version" based on input.
 60    """
 61    if isinstance(alias_or_version, int):
 62        return uri_for_model_version(name=name, version=alias_or_version)
 63    else:
 64        return uri_for_model_alias(name=name, alias=alias_or_version)
 65
 66
 67# %% SAVERS
 68
 69
 70class Saver(abc.ABC, pdt.BaseModel, strict=True, frozen=True, extra="forbid"):
 71    """Base class for saving models in registry.
 72
 73    Separate model definition from serialization.
 74    e.g., to switch between serialization flavors.
 75
 76    Parameters:
 77        path (str): model path inside the Mlflow store.
 78    """
 79
 80    KIND: str
 81
 82    path: str = "model"
 83
 84    @abc.abstractmethod
 85    def save(
 86        self,
 87        model: models.Model,
 88        signature: signers.Signature,
 89        input_example: schemas.Inputs,
 90    ) -> Info:
 91        """Save a model in the model registry.
 92
 93        Args:
 94            model (models.Model): project model to save.
 95            signature (signers.Signature): model signature.
 96            input_example (schemas.Inputs): sample of inputs.
 97
 98        Returns:
 99            Info: model saving information.
100        """
101
102
103class CustomSaver(Saver):
104    """Saver for project models using the Mlflow PyFunc module.
105
106    https://mlflow.org/docs/latest/python_api/mlflow.pyfunc.html
107    """
108
109    KIND: T.Literal["CustomSaver"] = "CustomSaver"
110
111    class Adapter(PythonModel):  # type: ignore[misc]
112        """Adapt a custom model to the Mlflow PyFunc flavor for saving operations.
113
114        https://mlflow.org/docs/latest/python_api/mlflow.pyfunc.html?#mlflow.pyfunc.PythonModel
115        """
116
117        def __init__(self, model: models.Model):
118            """Initialize the custom saver adapter.
119
120            Args:
121                model (models.Model): project model.
122            """
123            self.model = model
124
125        def predict(
126            self,
127            context: PythonModelContext,
128            model_input: schemas.Inputs,
129            params: dict[str, T.Any] | None = None,
130        ) -> schemas.Outputs:
131            """Generate predictions with a custom model for the given inputs.
132
133            Args:
134                context (mlflow.PythonModelContext): mlflow context.
135                model_input (schemas.Inputs): inputs for the mlflow model.
136                params (dict[str, T.Any] | None): additional parameters.
137
138            Returns:
139                schemas.Outputs: validated outputs of the project model.
140            """
141            return self.model.predict(inputs=model_input)
142
143    @T.override
144    def save(
145        self,
146        model: models.Model,
147        signature: signers.Signature,
148        input_example: schemas.Inputs,
149    ) -> Info:
150        adapter = CustomSaver.Adapter(model=model)
151        return mlflow.pyfunc.log_model(
152            python_model=adapter,
153            signature=signature,
154            artifact_path=self.path,
155            input_example=input_example,
156        )
157
158
159class BuiltinSaver(Saver):
160    """Saver for built-in models using an Mlflow flavor module.
161
162    https://mlflow.org/docs/latest/models.html#built-in-model-flavors
163
164    Parameters:
165        flavor (str): Mlflow flavor module to use for the serialization.
166    """
167
168    KIND: T.Literal["BuiltinSaver"] = "BuiltinSaver"
169
170    flavor: str
171
172    @T.override
173    def save(
174        self,
175        model: models.Model,
176        signature: signers.Signature,
177        input_example: schemas.Inputs,
178    ) -> Info:
179        builtin_model = model.get_internal_model()
180        module = getattr(mlflow, self.flavor)
181        return module.log_model(
182            builtin_model,
183            artifact_path=self.path,
184            signature=signature,
185            input_example=input_example,
186        )
187
188
189SaverKind = CustomSaver | BuiltinSaver
190
191# %% LOADERS
192
193
194class Loader(abc.ABC, pdt.BaseModel, strict=True, frozen=True, extra="forbid"):
195    """Base class for loading models from registry.
196
197    Separate model definition from deserialization.
198    e.g., to switch between deserialization flavors.
199    """
200
201    KIND: str
202
203    class Adapter(abc.ABC):
204        """Adapt any model for the project inference."""
205
206        @abc.abstractmethod
207        def predict(self, inputs: schemas.Inputs) -> schemas.Outputs:
208            """Generate predictions with the internal model for the given inputs.
209
210            Args:
211                inputs (schemas.Inputs): validated inputs for the project model.
212
213            Returns:
214                schemas.Outputs: validated outputs of the project model.
215            """
216
217    @abc.abstractmethod
218    def load(self, uri: str) -> "Loader.Adapter":
219        """Load a model from the model registry.
220
221        Args:
222            uri (str): URI of a model to load.
223
224        Returns:
225            Loader.Adapter: model loaded.
226        """
227
228
229class CustomLoader(Loader):
230    """Loader for custom models using the Mlflow PyFunc module.
231
232    https://mlflow.org/docs/latest/python_api/mlflow.pyfunc.html
233    """
234
235    KIND: T.Literal["CustomLoader"] = "CustomLoader"
236
237    class Adapter(Loader.Adapter):
238        """Adapt a custom model for the project inference."""
239
240        def __init__(self, model: PyFuncModel) -> None:
241            """Initialize the adapter from an mlflow pyfunc model.
242
243            Args:
244                model (PyFuncModel): mlflow pyfunc model.
245            """
246            self.model = model
247
248        @T.override
249        def predict(self, inputs: schemas.Inputs) -> schemas.Outputs:
250            # model validation is already done in predict
251            outputs = self.model.predict(data=inputs)
252            return T.cast(schemas.Outputs, outputs)
253
254    @T.override
255    def load(self, uri: str) -> "CustomLoader.Adapter":
256        model = mlflow.pyfunc.load_model(model_uri=uri)
257        adapter = CustomLoader.Adapter(model=model)
258        return adapter
259
260
261class BuiltinLoader(Loader):
262    """Loader for built-in models using the Mlflow PyFunc module.
263
264    Note: use Mlflow PyFunc instead of flavors to use standard API.
265
266    https://mlflow.org/docs/latest/models.html#built-in-model-flavors
267    """
268
269    KIND: T.Literal["BuiltinLoader"] = "BuiltinLoader"
270
271    class Adapter(Loader.Adapter):
272        """Adapt a builtin model for the project inference."""
273
274        def __init__(self, model: PyFuncModel) -> None:
275            """Initialize the adapter from an mlflow pyfunc model.
276
277            Args:
278                model (PyFuncModel): mlflow pyfunc model.
279            """
280            self.model = model
281
282        @T.override
283        def predict(self, inputs: schemas.Inputs) -> schemas.Outputs:
284            columns = list(schemas.OutputsSchema.to_schema().columns)
285            outputs = self.model.predict(data=inputs)  # unchecked data!
286            return schemas.Outputs(outputs, columns=columns, index=inputs.index)
287
288    @T.override
289    def load(self, uri: str) -> "BuiltinLoader.Adapter":
290        model = mlflow.pyfunc.load_model(model_uri=uri)
291        adapter = BuiltinLoader.Adapter(model=model)
292        return adapter
293
294
295LoaderKind = CustomLoader | BuiltinLoader
296
297# %% REGISTERS
298
299
300class Register(abc.ABC, pdt.BaseModel, strict=True, frozen=True, extra="forbid"):
301    """Base class for registring models to a location.
302
303    Separate model definition from its registration.
304    e.g., to change the model registry backend.
305
306    Parameters:
307        tags (dict[str, T.Any]): tags for the model.
308    """
309
310    KIND: str
311
312    tags: dict[str, T.Any] = {}
313
314    @abc.abstractmethod
315    def register(self, name: str, model_uri: str) -> Version:
316        """Register a model given its name and URI.
317
318        Args:
319            name (str): name of the model to register.
320            model_uri (str): URI of a model to register.
321
322        Returns:
323            Version: information about the registered model.
324        """
325
326
327class MlflowRegister(Register):
328    """Register for models in the Mlflow Model Registry.
329
330    https://mlflow.org/docs/latest/model-registry.html
331    """
332
333    KIND: T.Literal["MlflowRegister"] = "MlflowRegister"
334
335    @T.override
336    def register(self, name: str, model_uri: str) -> Version:
337        return mlflow.register_model(name=name, model_uri=model_uri, tags=self.tags)
338
339
340RegisterKind = MlflowRegister
Info: TypeAlias = mlflow.models.model.ModelInfo
Alias: TypeAlias = mlflow.entities.model_registry.model_version.ModelVersion
Version: TypeAlias = mlflow.entities.model_registry.model_version.ModelVersion
def uri_for_model_alias(name: str, alias: str) -> str: 
26def uri_for_model_alias(name: str, alias: str) -> str:
27    """Create a model URI from a model name and an alias.
28
29    Args:
30        name (str): name of the mlflow registered model.
31        alias (str): alias of the registered model.
32
33    Returns:
34        str: model URI as "models:/name@alias".
35    """
36    return f"models:/{name}@{alias}"

Create a model URI from a model name and an alias.

Arguments:
  • name (str): name of the mlflow registered model.
  • alias (str): alias of the registered model.
Returns:

str: model URI as "models:/name@alias".

def uri_for_model_version(name: str, version: int) -> str: 
39def uri_for_model_version(name: str, version: int) -> str:
40    """Create a model URI from a model name and a version.
41
42    Args:
43        name (str): name of the mlflow registered model.
44        version (int): version of the registered model.
45
46    Returns:
47        str: model URI as "models:/name/version."
48    """
49    return f"models:/{name}/{version}"

Create a model URI from a model name and a version.

Arguments:
  • name (str): name of the mlflow registered model.
  • version (int): version of the registered model.
Returns:

str: model URI as "models:/name/version."

def uri_for_model_alias_or_version(name: str, alias_or_version: str | int) -> str: 
52def uri_for_model_alias_or_version(name: str, alias_or_version: str | int) -> str:
53    """Create a model URi from a model name and an alias or version.
54
55    Args:
56        name (str): name of the mlflow registered model.
57        alias_or_version (str | int): alias or version of the registered model.
58
59    Returns:
60        str: model URI as "models:/name@alias" or "models:/name/version" based on input.
61    """
62    if isinstance(alias_or_version, int):
63        return uri_for_model_version(name=name, version=alias_or_version)
64    else:
65        return uri_for_model_alias(name=name, alias=alias_or_version)

Create a model URi from a model name and an alias or version.

Arguments:
  • name (str): name of the mlflow registered model.
  • alias_or_version (str | int): alias or version of the registered model.
Returns:

str: model URI as "models:/name@alias" or "models:/name/version" based on input.

class Saver(abc.ABC, pydantic.main.BaseModel): 
 71class Saver(abc.ABC, pdt.BaseModel, strict=True, frozen=True, extra="forbid"):
 72    """Base class for saving models in registry.
 73
 74    Separate model definition from serialization.
 75    e.g., to switch between serialization flavors.
 76
 77    Parameters:
 78        path (str): model path inside the Mlflow store.
 79    """
 80
 81    KIND: str
 82
 83    path: str = "model"
 84
 85    @abc.abstractmethod
 86    def save(
 87        self,
 88        model: models.Model,
 89        signature: signers.Signature,
 90        input_example: schemas.Inputs,
 91    ) -> Info:
 92        """Save a model in the model registry.
 93
 94        Args:
 95            model (models.Model): project model to save.
 96            signature (signers.Signature): model signature.
 97            input_example (schemas.Inputs): sample of inputs.
 98
 99        Returns:
100            Info: model saving information.
101        """

Base class for saving models in registry.

Separate model definition from serialization. e.g., to switch between serialization flavors.

Arguments:
  • path (str): model path inside the Mlflow store.
KIND: str
path: str
@abc.abstractmethod
def save( self, model: bikes.core.models.Model, signature: mlflow.models.signature.ModelSignature, input_example: pandera.typing.pandas.DataFrame[bikes.core.schemas.InputsSchema]) -> mlflow.models.model.ModelInfo: 
 85    @abc.abstractmethod
 86    def save(
 87        self,
 88        model: models.Model,
 89        signature: signers.Signature,
 90        input_example: schemas.Inputs,
 91    ) -> Info:
 92        """Save a model in the model registry.
 93
 94        Args:
 95            model (models.Model): project model to save.
 96            signature (signers.Signature): model signature.
 97            input_example (schemas.Inputs): sample of inputs.
 98
 99        Returns:
100            Info: model saving information.
101        """

Save a model in the model registry.

Arguments:
  • model (models.Model): project model to save.
  • signature (signers.Signature): model signature.
  • input_example (schemas.Inputs): sample of inputs.
Returns:

Info: model saving information.

model_config: ClassVar[pydantic.config.ConfigDict] = {'strict': True, 'frozen': True, 'extra': 'forbid'}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class CustomSaver(Saver): 
104class CustomSaver(Saver):
105    """Saver for project models using the Mlflow PyFunc module.
106
107    https://mlflow.org/docs/latest/python_api/mlflow.pyfunc.html
108    """
109
110    KIND: T.Literal["CustomSaver"] = "CustomSaver"
111
112    class Adapter(PythonModel):  # type: ignore[misc]
113        """Adapt a custom model to the Mlflow PyFunc flavor for saving operations.
114
115        https://mlflow.org/docs/latest/python_api/mlflow.pyfunc.html?#mlflow.pyfunc.PythonModel
116        """
117
118        def __init__(self, model: models.Model):
119            """Initialize the custom saver adapter.
120
121            Args:
122                model (models.Model): project model.
123            """
124            self.model = model
125
126        def predict(
127            self,
128            context: PythonModelContext,
129            model_input: schemas.Inputs,
130            params: dict[str, T.Any] | None = None,
131        ) -> schemas.Outputs:
132            """Generate predictions with a custom model for the given inputs.
133
134            Args:
135                context (mlflow.PythonModelContext): mlflow context.
136                model_input (schemas.Inputs): inputs for the mlflow model.
137                params (dict[str, T.Any] | None): additional parameters.
138
139            Returns:
140                schemas.Outputs: validated outputs of the project model.
141            """
142            return self.model.predict(inputs=model_input)
143
144    @T.override
145    def save(
146        self,
147        model: models.Model,
148        signature: signers.Signature,
149        input_example: schemas.Inputs,
150    ) -> Info:
151        adapter = CustomSaver.Adapter(model=model)
152        return mlflow.pyfunc.log_model(
153            python_model=adapter,
154            signature=signature,
155            artifact_path=self.path,
156            input_example=input_example,
157        )

Saver for project models using the Mlflow PyFunc module.

https://mlflow.org/docs/latest/python_api/mlflow.pyfunc.html

KIND: Literal['CustomSaver']
@T.override
def save( self, model: bikes.core.models.Model, signature: mlflow.models.signature.ModelSignature, input_example: pandera.typing.pandas.DataFrame[bikes.core.schemas.InputsSchema]) -> mlflow.models.model.ModelInfo: 
144    @T.override
145    def save(
146        self,
147        model: models.Model,
148        signature: signers.Signature,
149        input_example: schemas.Inputs,
150    ) -> Info:
151        adapter = CustomSaver.Adapter(model=model)
152        return mlflow.pyfunc.log_model(
153            python_model=adapter,
154            signature=signature,
155            artifact_path=self.path,
156            input_example=input_example,
157        )

Save a model in the model registry.

Arguments:
  • model (models.Model): project model to save.
  • signature (signers.Signature): model signature.
  • input_example (schemas.Inputs): sample of inputs.
Returns:

Info: model saving information.

model_config: ClassVar[pydantic.config.ConfigDict] = {'strict': True, 'frozen': True, 'extra': 'forbid'}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

Inherited Members
Saver
path
class CustomSaver.Adapter(mlflow.pyfunc.model.PythonModel): 
112    class Adapter(PythonModel):  # type: ignore[misc]
113        """Adapt a custom model to the Mlflow PyFunc flavor for saving operations.
114
115        https://mlflow.org/docs/latest/python_api/mlflow.pyfunc.html?#mlflow.pyfunc.PythonModel
116        """
117
118        def __init__(self, model: models.Model):
119            """Initialize the custom saver adapter.
120
121            Args:
122                model (models.Model): project model.
123            """
124            self.model = model
125
126        def predict(
127            self,
128            context: PythonModelContext,
129            model_input: schemas.Inputs,
130            params: dict[str, T.Any] | None = None,
131        ) -> schemas.Outputs:
132            """Generate predictions with a custom model for the given inputs.
133
134            Args:
135                context (mlflow.PythonModelContext): mlflow context.
136                model_input (schemas.Inputs): inputs for the mlflow model.
137                params (dict[str, T.Any] | None): additional parameters.
138
139            Returns:
140                schemas.Outputs: validated outputs of the project model.
141            """
142            return self.model.predict(inputs=model_input)

Adapt a custom model to the Mlflow PyFunc flavor for saving operations.

https://mlflow.org/docs/latest/python_api/mlflow.pyfunc.html?#mlflow.pyfunc.PythonModel

CustomSaver.Adapter(model: bikes.core.models.Model) 
118        def __init__(self, model: models.Model):
119            """Initialize the custom saver adapter.
120
121            Args:
122                model (models.Model): project model.
123            """
124            self.model = model

Initialize the custom saver adapter.

Arguments:
  • model (models.Model): project model.
model
def predict( self, context: mlflow.pyfunc.model.PythonModelContext, model_input: pandera.typing.pandas.DataFrame[bikes.core.schemas.InputsSchema], params: dict[str, typing.Any] | None = None) -> pandera.typing.pandas.DataFrame[bikes.core.schemas.OutputsSchema]: 
126        def predict(
127            self,
128            context: PythonModelContext,
129            model_input: schemas.Inputs,
130            params: dict[str, T.Any] | None = None,
131        ) -> schemas.Outputs:
132            """Generate predictions with a custom model for the given inputs.
133
134            Args:
135                context (mlflow.PythonModelContext): mlflow context.
136                model_input (schemas.Inputs): inputs for the mlflow model.
137                params (dict[str, T.Any] | None): additional parameters.
138
139            Returns:
140                schemas.Outputs: validated outputs of the project model.
141            """
142            return self.model.predict(inputs=model_input)

Generate predictions with a custom model for the given inputs.

Arguments:
  • context (mlflow.PythonModelContext): mlflow context.
  • model_input (schemas.Inputs): inputs for the mlflow model.
  • params (dict[str, T.Any] | None): additional parameters.
Returns:

schemas.Outputs: validated outputs of the project model.

class BuiltinSaver(Saver): 
160class BuiltinSaver(Saver):
161    """Saver for built-in models using an Mlflow flavor module.
162
163    https://mlflow.org/docs/latest/models.html#built-in-model-flavors
164
165    Parameters:
166        flavor (str): Mlflow flavor module to use for the serialization.
167    """
168
169    KIND: T.Literal["BuiltinSaver"] = "BuiltinSaver"
170
171    flavor: str
172
173    @T.override
174    def save(
175        self,
176        model: models.Model,
177        signature: signers.Signature,
178        input_example: schemas.Inputs,
179    ) -> Info:
180        builtin_model = model.get_internal_model()
181        module = getattr(mlflow, self.flavor)
182        return module.log_model(
183            builtin_model,
184            artifact_path=self.path,
185            signature=signature,
186            input_example=input_example,
187        )

Saver for built-in models using an Mlflow flavor module.

https://mlflow.org/docs/latest/models.html#built-in-model-flavors

Arguments:
  • flavor (str): Mlflow flavor module to use for the serialization.
KIND: Literal['BuiltinSaver']
flavor: str
@T.override
def save( self, model: bikes.core.models.Model, signature: mlflow.models.signature.ModelSignature, input_example: pandera.typing.pandas.DataFrame[bikes.core.schemas.InputsSchema]) -> mlflow.models.model.ModelInfo: 
173    @T.override
174    def save(
175        self,
176        model: models.Model,
177        signature: signers.Signature,
178        input_example: schemas.Inputs,
179    ) -> Info:
180        builtin_model = model.get_internal_model()
181        module = getattr(mlflow, self.flavor)
182        return module.log_model(
183            builtin_model,
184            artifact_path=self.path,
185            signature=signature,
186            input_example=input_example,
187        )

Save a model in the model registry.

Arguments:
  • model (models.Model): project model to save.
  • signature (signers.Signature): model signature.
  • input_example (schemas.Inputs): sample of inputs.
Returns:

Info: model saving information.

model_config: ClassVar[pydantic.config.ConfigDict] = {'strict': True, 'frozen': True, 'extra': 'forbid'}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

Inherited Members
Saver
path
SaverKind = CustomSaver | BuiltinSaver
class Loader(abc.ABC, pydantic.main.BaseModel): 
195class Loader(abc.ABC, pdt.BaseModel, strict=True, frozen=True, extra="forbid"):
196    """Base class for loading models from registry.
197
198    Separate model definition from deserialization.
199    e.g., to switch between deserialization flavors.
200    """
201
202    KIND: str
203
204    class Adapter(abc.ABC):
205        """Adapt any model for the project inference."""
206
207        @abc.abstractmethod
208        def predict(self, inputs: schemas.Inputs) -> schemas.Outputs:
209            """Generate predictions with the internal model for the given inputs.
210
211            Args:
212                inputs (schemas.Inputs): validated inputs for the project model.
213
214            Returns:
215                schemas.Outputs: validated outputs of the project model.
216            """
217
218    @abc.abstractmethod
219    def load(self, uri: str) -> "Loader.Adapter":
220        """Load a model from the model registry.
221
222        Args:
223            uri (str): URI of a model to load.
224
225        Returns:
226            Loader.Adapter: model loaded.
227        """

Base class for loading models from registry.

Separate model definition from deserialization. e.g., to switch between deserialization flavors.

KIND: str
@abc.abstractmethod
def load(self, uri: str) -> Loader.Adapter: 
218    @abc.abstractmethod
219    def load(self, uri: str) -> "Loader.Adapter":
220        """Load a model from the model registry.
221
222        Args:
223            uri (str): URI of a model to load.
224
225        Returns:
226            Loader.Adapter: model loaded.
227        """

Load a model from the model registry.

Arguments:
  • uri (str): URI of a model to load.
Returns:

Loader.Adapter: model loaded.

model_config: ClassVar[pydantic.config.ConfigDict] = {'strict': True, 'frozen': True, 'extra': 'forbid'}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class Loader.Adapter(abc.ABC): 
204    class Adapter(abc.ABC):
205        """Adapt any model for the project inference."""
206
207        @abc.abstractmethod
208        def predict(self, inputs: schemas.Inputs) -> schemas.Outputs:
209            """Generate predictions with the internal model for the given inputs.
210
211            Args:
212                inputs (schemas.Inputs): validated inputs for the project model.
213
214            Returns:
215                schemas.Outputs: validated outputs of the project model.
216            """

Adapt any model for the project inference.

@abc.abstractmethod
def predict( self, inputs: pandera.typing.pandas.DataFrame[bikes.core.schemas.InputsSchema]) -> pandera.typing.pandas.DataFrame[bikes.core.schemas.OutputsSchema]: 
207        @abc.abstractmethod
208        def predict(self, inputs: schemas.Inputs) -> schemas.Outputs:
209            """Generate predictions with the internal model for the given inputs.
210
211            Args:
212                inputs (schemas.Inputs): validated inputs for the project model.
213
214            Returns:
215                schemas.Outputs: validated outputs of the project model.
216            """

Generate predictions with the internal model for the given inputs.

Arguments:
  • inputs (schemas.Inputs): validated inputs for the project model.
Returns:

schemas.Outputs: validated outputs of the project model.

class CustomLoader(Loader): 
230class CustomLoader(Loader):
231    """Loader for custom models using the Mlflow PyFunc module.
232
233    https://mlflow.org/docs/latest/python_api/mlflow.pyfunc.html
234    """
235
236    KIND: T.Literal["CustomLoader"] = "CustomLoader"
237
238    class Adapter(Loader.Adapter):
239        """Adapt a custom model for the project inference."""
240
241        def __init__(self, model: PyFuncModel) -> None:
242            """Initialize the adapter from an mlflow pyfunc model.
243
244            Args:
245                model (PyFuncModel): mlflow pyfunc model.
246            """
247            self.model = model
248
249        @T.override
250        def predict(self, inputs: schemas.Inputs) -> schemas.Outputs:
251            # model validation is already done in predict
252            outputs = self.model.predict(data=inputs)
253            return T.cast(schemas.Outputs, outputs)
254
255    @T.override
256    def load(self, uri: str) -> "CustomLoader.Adapter":
257        model = mlflow.pyfunc.load_model(model_uri=uri)
258        adapter = CustomLoader.Adapter(model=model)
259        return adapter

Loader for custom models using the Mlflow PyFunc module.

https://mlflow.org/docs/latest/python_api/mlflow.pyfunc.html

KIND: Literal['CustomLoader']
@T.override
def load(self, uri: str) -> CustomLoader.Adapter: 
255    @T.override
256    def load(self, uri: str) -> "CustomLoader.Adapter":
257        model = mlflow.pyfunc.load_model(model_uri=uri)
258        adapter = CustomLoader.Adapter(model=model)
259        return adapter

Load a model from the model registry.

Arguments:
  • uri (str): URI of a model to load.
Returns:

Loader.Adapter: model loaded.

model_config: ClassVar[pydantic.config.ConfigDict] = {'strict': True, 'frozen': True, 'extra': 'forbid'}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class CustomLoader.Adapter(Loader.Adapter): 
238    class Adapter(Loader.Adapter):
239        """Adapt a custom model for the project inference."""
240
241        def __init__(self, model: PyFuncModel) -> None:
242            """Initialize the adapter from an mlflow pyfunc model.
243
244            Args:
245                model (PyFuncModel): mlflow pyfunc model.
246            """
247            self.model = model
248
249        @T.override
250        def predict(self, inputs: schemas.Inputs) -> schemas.Outputs:
251            # model validation is already done in predict
252            outputs = self.model.predict(data=inputs)
253            return T.cast(schemas.Outputs, outputs)

Adapt a custom model for the project inference.

CustomLoader.Adapter(model: mlflow.pyfunc.PyFuncModel) 
241        def __init__(self, model: PyFuncModel) -> None:
242            """Initialize the adapter from an mlflow pyfunc model.
243
244            Args:
245                model (PyFuncModel): mlflow pyfunc model.
246            """
247            self.model = model

Initialize the adapter from an mlflow pyfunc model.

Arguments:
  • model (PyFuncModel): mlflow pyfunc model.
model
@T.override
def predict( self, inputs: pandera.typing.pandas.DataFrame[bikes.core.schemas.InputsSchema]) -> pandera.typing.pandas.DataFrame[bikes.core.schemas.OutputsSchema]: 
249        @T.override
250        def predict(self, inputs: schemas.Inputs) -> schemas.Outputs:
251            # model validation is already done in predict
252            outputs = self.model.predict(data=inputs)
253            return T.cast(schemas.Outputs, outputs)

Generate predictions with the internal model for the given inputs.

Arguments:
  • inputs (schemas.Inputs): validated inputs for the project model.
Returns:

schemas.Outputs: validated outputs of the project model.

class BuiltinLoader(Loader): 
262class BuiltinLoader(Loader):
263    """Loader for built-in models using the Mlflow PyFunc module.
264
265    Note: use Mlflow PyFunc instead of flavors to use standard API.
266
267    https://mlflow.org/docs/latest/models.html#built-in-model-flavors
268    """
269
270    KIND: T.Literal["BuiltinLoader"] = "BuiltinLoader"
271
272    class Adapter(Loader.Adapter):
273        """Adapt a builtin model for the project inference."""
274
275        def __init__(self, model: PyFuncModel) -> None:
276            """Initialize the adapter from an mlflow pyfunc model.
277
278            Args:
279                model (PyFuncModel): mlflow pyfunc model.
280            """
281            self.model = model
282
283        @T.override
284        def predict(self, inputs: schemas.Inputs) -> schemas.Outputs:
285            columns = list(schemas.OutputsSchema.to_schema().columns)
286            outputs = self.model.predict(data=inputs)  # unchecked data!
287            return schemas.Outputs(outputs, columns=columns, index=inputs.index)
288
289    @T.override
290    def load(self, uri: str) -> "BuiltinLoader.Adapter":
291        model = mlflow.pyfunc.load_model(model_uri=uri)
292        adapter = BuiltinLoader.Adapter(model=model)
293        return adapter

Loader for built-in models using the Mlflow PyFunc module.

Note: use Mlflow PyFunc instead of flavors to use standard API.

https://mlflow.org/docs/latest/models.html#built-in-model-flavors

KIND: Literal['BuiltinLoader']
@T.override
def load(self, uri: str) -> BuiltinLoader.Adapter: 
289    @T.override
290    def load(self, uri: str) -> "BuiltinLoader.Adapter":
291        model = mlflow.pyfunc.load_model(model_uri=uri)
292        adapter = BuiltinLoader.Adapter(model=model)
293        return adapter

Load a model from the model registry.

Arguments:
  • uri (str): URI of a model to load.
Returns:

Loader.Adapter: model loaded.

model_config: ClassVar[pydantic.config.ConfigDict] = {'strict': True, 'frozen': True, 'extra': 'forbid'}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class BuiltinLoader.Adapter(Loader.Adapter): 
272    class Adapter(Loader.Adapter):
273        """Adapt a builtin model for the project inference."""
274
275        def __init__(self, model: PyFuncModel) -> None:
276            """Initialize the adapter from an mlflow pyfunc model.
277
278            Args:
279                model (PyFuncModel): mlflow pyfunc model.
280            """
281            self.model = model
282
283        @T.override
284        def predict(self, inputs: schemas.Inputs) -> schemas.Outputs:
285            columns = list(schemas.OutputsSchema.to_schema().columns)
286            outputs = self.model.predict(data=inputs)  # unchecked data!
287            return schemas.Outputs(outputs, columns=columns, index=inputs.index)

Adapt a builtin model for the project inference.

BuiltinLoader.Adapter(model: mlflow.pyfunc.PyFuncModel) 
275        def __init__(self, model: PyFuncModel) -> None:
276            """Initialize the adapter from an mlflow pyfunc model.
277
278            Args:
279                model (PyFuncModel): mlflow pyfunc model.
280            """
281            self.model = model

Initialize the adapter from an mlflow pyfunc model.

Arguments:
  • model (PyFuncModel): mlflow pyfunc model.
model
@T.override
def predict( self, inputs: pandera.typing.pandas.DataFrame[bikes.core.schemas.InputsSchema]) -> pandera.typing.pandas.DataFrame[bikes.core.schemas.OutputsSchema]: 
283        @T.override
284        def predict(self, inputs: schemas.Inputs) -> schemas.Outputs:
285            columns = list(schemas.OutputsSchema.to_schema().columns)
286            outputs = self.model.predict(data=inputs)  # unchecked data!
287            return schemas.Outputs(outputs, columns=columns, index=inputs.index)

Generate predictions with the internal model for the given inputs.

Arguments:
  • inputs (schemas.Inputs): validated inputs for the project model.
Returns:

schemas.Outputs: validated outputs of the project model.

LoaderKind = CustomLoader | BuiltinLoader
class Register(abc.ABC, pydantic.main.BaseModel): 
301class Register(abc.ABC, pdt.BaseModel, strict=True, frozen=True, extra="forbid"):
302    """Base class for registring models to a location.
303
304    Separate model definition from its registration.
305    e.g., to change the model registry backend.
306
307    Parameters:
308        tags (dict[str, T.Any]): tags for the model.
309    """
310
311    KIND: str
312
313    tags: dict[str, T.Any] = {}
314
315    @abc.abstractmethod
316    def register(self, name: str, model_uri: str) -> Version:
317        """Register a model given its name and URI.
318
319        Args:
320            name (str): name of the model to register.
321            model_uri (str): URI of a model to register.
322
323        Returns:
324            Version: information about the registered model.
325        """

Base class for registring models to a location.

Separate model definition from its registration. e.g., to change the model registry backend.

Arguments:
  • tags (dict[str, T.Any]): tags for the model.
KIND: str
tags: dict[str, typing.Any]
@abc.abstractmethod
def register( self, name: str, model_uri: str) -> mlflow.entities.model_registry.model_version.ModelVersion: 
315    @abc.abstractmethod
316    def register(self, name: str, model_uri: str) -> Version:
317        """Register a model given its name and URI.
318
319        Args:
320            name (str): name of the model to register.
321            model_uri (str): URI of a model to register.
322
323        Returns:
324            Version: information about the registered model.
325        """

Register a model given its name and URI.

Arguments:
  • name (str): name of the model to register.
  • model_uri (str): URI of a model to register.
Returns:

Version: information about the registered model.

model_config: ClassVar[pydantic.config.ConfigDict] = {'strict': True, 'frozen': True, 'extra': 'forbid'}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class MlflowRegister(Register): 
328class MlflowRegister(Register):
329    """Register for models in the Mlflow Model Registry.
330
331    https://mlflow.org/docs/latest/model-registry.html
332    """
333
334    KIND: T.Literal["MlflowRegister"] = "MlflowRegister"
335
336    @T.override
337    def register(self, name: str, model_uri: str) -> Version:
338        return mlflow.register_model(name=name, model_uri=model_uri, tags=self.tags)

Register for models in the Mlflow Model Registry.

https://mlflow.org/docs/latest/model-registry.html

KIND: Literal['MlflowRegister']
@T.override
def register( self, name: str, model_uri: str) -> mlflow.entities.model_registry.model_version.ModelVersion: 
336    @T.override
337    def register(self, name: str, model_uri: str) -> Version:
338        return mlflow.register_model(name=name, model_uri=model_uri, tags=self.tags)

Register a model given its name and URI.

Arguments:
  • name (str): name of the model to register.
  • model_uri (str): URI of a model to register.
Returns:

Version: information about the registered model.

model_config: ClassVar[pydantic.config.ConfigDict] = {'strict': True, 'frozen': True, 'extra': 'forbid'}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

Inherited Members
Register
tags
RegisterKind = <class 'MlflowRegister'>