1  import os
 2  from typing import Any
 3  
 4  from custom_code import iris_classes
 5  from sklearn.datasets import load_iris
 6  from sklearn.linear_model import LogisticRegression
 7  
 8  import mlflow
 9  from mlflow.models import infer_signature
10  
11  
12  class CustomPredict(mlflow.pyfunc.PythonModel):
13      """Custom pyfunc class used to create customized mlflow models"""
14  
15      def load_context(self, context):
16          self.model = mlflow.sklearn.load_model(context.artifacts["custom_model"])
17  
18      def predict(self, context, model_input, params: dict[str, Any] | None = None):
19          prediction = self.model.predict(model_input)
20          return iris_classes(prediction)
21  
22  
23  X, y = load_iris(return_X_y=True, as_frame=True)
24  params = {"C": 1.0, "random_state": 42}
25  classifier = LogisticRegression(**params).fit(X, y)
26  
27  predictions = classifier.predict(X)
28  signature = infer_signature(X, predictions)
29  
30  with mlflow.start_run(run_name="test_pyfunc") as run:
31      model_info = mlflow.sklearn.log_model(sk_model=classifier, name="model", signature=signature)
32  
33      # start a child run to create custom imagine model
34      with mlflow.start_run(run_name="test_custom_model", nested=True):
35          print(f"Pyfunc run ID: {run.info.run_id}")
36          # log a custom model
37          mlflow.pyfunc.log_model(
38              name="artifacts",
39              code_paths=[os.getcwd()],
40              artifacts={"custom_model": model_info.model_uri},
41              python_model=CustomPredict(),
42              signature=signature,
43          )