1  from typing import Any
 2  
 3  from custom_code import iris_classes
 4  
 5  import mlflow
 6  
 7  
 8  class CustomPredict(mlflow.pyfunc.PythonModel):
 9      """Custom pyfunc class used to create customized mlflow models"""
10  
11      def predict(self, context, model_input, params: dict[str, Any] | None = None):
12          prediction = [x % 3 for x in model_input]
13          return iris_classes(prediction)
14  
15  
16  with mlflow.start_run(run_name="test_custom_model_with_inferred_code_paths"):
17      # log a custom model
18      model_info = mlflow.pyfunc.log_model(
19          name="artifacts",
20          infer_code_paths=True,
21          python_model=CustomPredict(),
22      )
23      print(f"Model URI: {model_info.model_uri}")