chore: import upstream snapshot with attribution

examples/pyfunc/README.md

@@ -0,0 +1,53 @@
+# Pyfunc model example
+
+This example demonstrates the use of a pyfunc model with custom inference logic.
+More specifically:
+
+- train a simple classification model
+- create a _pyfunc_ model that encapsulates the classification model with an attached module for custom inference logic
+
+## Structure of this example
+
+This examples contains a `train.py` file that trains a scikit-learn model with iris dataset and uses MLflow Tracking APIs to log the model. The nested **mlflow run** delivers the packaging of `pyfunc` model and `custom_code` module is attached
+to act as a custom inference logic layer in inference time.
+
+```
+├── train.py
+├── infer_model_code_path.py
+└── custom_code.py
+```
+
+## Running this example
+
+1. Train and log the model
+
+```
+$ python train.py
+```
+
+or train and log the model using inferred code paths
+
+```
+$ python infer_model_code_paths.py
+```
+
+2. Serve the pyfunc model
+
+```bash
+# Replace <pyfunc_run_id> with the run ID obtained in the previous step
+$ mlflow models serve -m "runs:/<pyfunc_run_id>/model" -p 5001
+```
+
+3. Send a request
+
+```
+$ curl http://127.0.0.1:5001/invocations -H 'Content-Type: application/json' -d '{
+  "dataframe_records": [[1, 1, 1, 1]]
+}'
+```
+
+The response should look like this:
+
+```
+[0]
+```

examples/pyfunc/custom_code.py

@@ -0,0 +1,5 @@
+flower_classes = ["setosa", "versicolor", "virginica"]
+
+
+def iris_classes(preds):
+    return [flower_classes[x] for x in preds]

examples/pyfunc/infer_model_code_paths.py

@@ -0,0 +1,23 @@
+from typing import Any
+
+from custom_code import iris_classes
+
+import mlflow
+
+
+class CustomPredict(mlflow.pyfunc.PythonModel):
+    """Custom pyfunc class used to create customized mlflow models"""
+
+    def predict(self, context, model_input, params: dict[str, Any] | None = None):
+        prediction = [x % 3 for x in model_input]
+        return iris_classes(prediction)
+
+
+with mlflow.start_run(run_name="test_custom_model_with_inferred_code_paths"):
+    # log a custom model
+    model_info = mlflow.pyfunc.log_model(
+        name="artifacts",
+        infer_code_paths=True,
+        python_model=CustomPredict(),
+    )
+    print(f"Model URI: {model_info.model_uri}")

examples/pyfunc/model_as_code.py

@@ -0,0 +1,52 @@
+# This example demonstrates defining a model directly from code.
+# This feature allows for defining model logic within a python script, module, or notebook that is stored
+# directly as serialized code, as opposed to object serialization that would otherwise occur when saving
+# or logging a model object.
+# This script defines the model's logic and specifies which class within the file contains the model code.
+# The companion example to this, model_as_code_driver.py, is the driver code that performs the  logging and
+# loading of this model definition.
+import os
+
+import pandas as pd
+
+import mlflow
+from mlflow import pyfunc
+
+assert "OPENAI_API_KEY" in os.environ, "Please set the OPENAI_API_KEY environment variable."
+
+
+class AIModel(pyfunc.PythonModel):
+    @mlflow.trace(name="chain", span_type="CHAIN")
+    def predict(self, context, model_input):
+        if isinstance(model_input, pd.DataFrame):
+            model_input = model_input["input"].tolist()
+
+        responses = []
+        for user_input in model_input:
+            response = self.get_open_ai_model_response(str(user_input))
+            responses.append(response.choices[0].message.content)
+
+        return pd.DataFrame({"response": responses})
+
+    @mlflow.trace(name="open_ai", span_type="LLM")
+    def get_open_ai_model_response(self, user_input):
+        from openai import OpenAI
+
+        return OpenAI().chat.completions.create(
+            model="gpt-4o-mini",
+            messages=[
+                {
+                    "role": "system",
+                    "content": "You are a helpful assistant. You are here to provide useful information to the user.",
+                },
+                {
+                    "role": "user",
+                    "content": user_input,
+                },
+            ],
+        )
+
+
+# IMPORTANT: The model code needs to call `mlflow.models.set_model()` to set the model,
+# which will be loaded back using `mlflow.pyfunc.load_model` for inference.
+mlflow.models.set_model(AIModel())

examples/pyfunc/model_as_code_driver.py

@@ -0,0 +1,28 @@
+# This is an example for logging a Python model from code using the
+# mlflow.pyfunc.log_model API. When a path to a valid Python script is submitted to the
+# python_model argument, the model code itself is serialized instead of the model object.
+# Within the targeted script, the model implementation must be defined and set by
+# using the mlflow.models.set_model API.
+
+import pandas as pd
+
+import mlflow
+
+input_example = ["What is the weather like today?"]
+
+# Specify the path to the model notebook
+model_path = "model_as_code.py"
+print(f"Model path: {model_path}")
+
+print("Logging model as code using Pyfunc log model API")
+with mlflow.start_run():
+    model_info = mlflow.pyfunc.log_model(
+        python_model=model_path,
+        name="ai-model",
+        input_example=input_example,
+    )
+
+print("Loading model using Pyfunc load model API")
+pyfunc_model = mlflow.pyfunc.load_model(model_info.model_uri)
+output = pyfunc_model.predict(pd.DataFrame(input_example, columns=["input"]))
+print(f"Output: {output}")

examples/pyfunc/train.py

@@ -0,0 +1,43 @@
+import os
+from typing import Any
+
+from custom_code import iris_classes
+from sklearn.datasets import load_iris
+from sklearn.linear_model import LogisticRegression
+
+import mlflow
+from mlflow.models import infer_signature
+
+
+class CustomPredict(mlflow.pyfunc.PythonModel):
+    """Custom pyfunc class used to create customized mlflow models"""
+
+    def load_context(self, context):
+        self.model = mlflow.sklearn.load_model(context.artifacts["custom_model"])
+
+    def predict(self, context, model_input, params: dict[str, Any] | None = None):
+        prediction = self.model.predict(model_input)
+        return iris_classes(prediction)
+
+
+X, y = load_iris(return_X_y=True, as_frame=True)
+params = {"C": 1.0, "random_state": 42}
+classifier = LogisticRegression(**params).fit(X, y)
+
+predictions = classifier.predict(X)
+signature = infer_signature(X, predictions)
+
+with mlflow.start_run(run_name="test_pyfunc") as run:
+    model_info = mlflow.sklearn.log_model(sk_model=classifier, name="model", signature=signature)
+
+    # start a child run to create custom imagine model
+    with mlflow.start_run(run_name="test_custom_model", nested=True):
+        print(f"Pyfunc run ID: {run.info.run_id}")
+        # log a custom model
+        mlflow.pyfunc.log_model(
+            name="artifacts",
+            code_paths=[os.getcwd()],
+            artifacts={"custom_model": model_info.model_uri},
+            python_model=CustomPredict(),
+            signature=signature,
+        )