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

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 = PydanticUndefined
path: str = 'model'
@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: 
 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        """

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.

class CustomSaver(Saver): 
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,  # noqa: ARG002  # required by mlflow PythonModel.predict
128            model_input: schemas.Inputs,
129            params: dict[str, T.Any] | None = None,  # noqa: ARG002  # required by mlflow PythonModel.predict
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            name=self.path,
155            input_example=input_example,
156        )

Saver for project models using the Mlflow PyFunc module.

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

KIND: Literal['CustomSaver'] = '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: 
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            name=self.path,
155            input_example=input_example,
156        )

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.

Inherited Members
Saver
path
class CustomSaver.Adapter(mlflow.pyfunc.model.PythonModel): 
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,  # noqa: ARG002  # required by mlflow PythonModel.predict
128            model_input: schemas.Inputs,
129            params: dict[str, T.Any] | None = None,  # noqa: ARG002  # required by mlflow PythonModel.predict
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)

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

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, Any] | None = None) -> pandera.typing.pandas.DataFrame[bikes.core.schemas.OutputsSchema]: 
125        def predict(
126            self,
127            context: PythonModelContext,  # noqa: ARG002  # required by mlflow PythonModel.predict
128            model_input: schemas.Inputs,
129            params: dict[str, T.Any] | None = None,  # noqa: ARG002  # required by mlflow PythonModel.predict
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)

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): 
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            name=self.path,
184            signature=signature,
185            input_example=input_example,
186        )

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'] = 'BuiltinSaver'
flavor: str = PydanticUndefined
@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: 
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            name=self.path,
184            signature=signature,
185            input_example=input_example,
186        )

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.

Inherited Members
Saver
path
SaverKind = CustomSaver | BuiltinSaver
class Loader(abc.ABC, pydantic.main.BaseModel): 
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        """

Base class for loading models from registry.

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

KIND: str = PydanticUndefined
@abc.abstractmethod
def load(self, uri: str) -> Loader.Adapter: 
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        """

Load a model from the model registry.

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

Loader.Adapter: model loaded.

class Loader.Adapter(abc.ABC): 
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            """

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

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): 
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        return CustomLoader.Adapter(model=model)

Loader for custom models using the Mlflow PyFunc module.

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

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

Load a model from the model registry.

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

Loader.Adapter: model loaded.

class CustomLoader.Adapter(Loader.Adapter): 
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)

Adapt a custom model for the project inference.

CustomLoader.Adapter(model: mlflow.pyfunc.PyFuncModel) 
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

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

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

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'] = 'BuiltinLoader'
@T.override
def load(self, uri: str) -> BuiltinLoader.Adapter: 
287    @T.override
288    def load(self, uri: str) -> BuiltinLoader.Adapter:
289        model = mlflow.pyfunc.load_model(model_uri=uri)
290        return BuiltinLoader.Adapter(model=model)

Load a model from the model registry.

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

Loader.Adapter: model loaded.

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

Adapt a builtin model for the project inference.

BuiltinLoader.Adapter(model: mlflow.pyfunc.PyFuncModel) 
273        def __init__(self, model: PyFuncModel) -> None:
274            """Initialize the adapter from an mlflow pyfunc model.
275
276            Args:
277                model (PyFuncModel): mlflow pyfunc model.
278            """
279            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]: 
281        @T.override
282        def predict(self, inputs: schemas.Inputs) -> schemas.Outputs:
283            columns = list(schemas.OutputsSchema.to_schema().columns)
284            outputs = self.model.predict(data=inputs)  # unchecked data!
285            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): 
298class Register(abc.ABC, pdt.BaseModel, strict=True, frozen=True, extra="forbid"):
299    """Base class for registring models to a location.
300
301    Separate model definition from its registration.
302    e.g., to change the model registry backend.
303
304    Parameters:
305        tags (dict[str, T.Any]): tags for the model.
306    """
307
308    KIND: str
309
310    tags: dict[str, T.Any] = {}
311
312    @abc.abstractmethod
313    def register(self, name: str, model_uri: str) -> Version:
314        """Register a model given its name and URI.
315
316        Args:
317            name (str): name of the model to register.
318            model_uri (str): URI of a model to register.
319
320        Returns:
321            Version: information about the registered model.
322        """

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 = PydanticUndefined
tags: dict[str, typing.Any] = {}
@abc.abstractmethod
def register( self, name: str, model_uri: str) -> mlflow.entities.model_registry.model_version.ModelVersion: 
312    @abc.abstractmethod
313    def register(self, name: str, model_uri: str) -> Version:
314        """Register a model given its name and URI.
315
316        Args:
317            name (str): name of the model to register.
318            model_uri (str): URI of a model to register.
319
320        Returns:
321            Version: information about the registered model.
322        """

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.

class MlflowRegister(Register): 
325class MlflowRegister(Register):
326    """Register for models in the Mlflow Model Registry.
327
328    https://mlflow.org/docs/latest/model-registry.html
329    """
330
331    KIND: T.Literal["MlflowRegister"] = "MlflowRegister"
332
333    @T.override
334    def register(self, name: str, model_uri: str) -> Version:
335        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'] = 'MlflowRegister'
@T.override
def register( self, name: str, model_uri: str) -> mlflow.entities.model_registry.model_version.ModelVersion: 
333    @T.override
334    def register(self, name: str, model_uri: str) -> Version:
335        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.

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