chore: import upstream snapshot with attribution

This commit is contained in:
wehub-resource-sync
2026-07-13 13:22:34 +08:00
commit 4b22cfda96
9037 changed files with 2363717 additions and 0 deletions
+40
View File
@@ -0,0 +1,40 @@
## MLflow examples
### Quick Start example
- `quickstart/mlflow_tracking.py` is a basic example to introduce MLflow concepts.
## Tutorials
Various examples that depict MLflow tracking, project, and serving use cases.
- `h2o` depicts how MLflow can be use to track various random forest architectures to train models
for predicting wine quality.
- `hyperparam` shows how to do hyperparameter tuning with MLflow and some popular optimization libraries.
- `keras` modifies
[a Keras classification example](https://github.com/keras-team/keras/blob/ed07472bc5fc985982db355135d37059a1f887a9/examples/reuters_mlp.py)
and uses MLflow's `mlflow.tensorflow.autolog()` API to automatically log metrics and parameters
to MLflow during training.
- `multistep_workflow` is an end-to-end of a data ETL and ML training pipeline built as an MLflow
project. The example shows how parts of the workflow can leverage from previously run steps.
- `pytorch` uses CNN on MNIST dataset for character recognition. The example logs TensorBoard events
and stores (logs) them as MLflow artifacts.
- `remote_store` has a usage example of REST based backed store for tracking.
- `r_wine` demonstrates how to log parameters, metrics, and models from R.
- `sklearn_elasticnet_diabetes` uses the sklearn diabetes dataset to predict diabetes progression
using ElasticNet.
- `sklearn_elasticnet_wine_quality` is an example for MLflow projects. This uses the Wine
Quality dataset and Elastic Net to predict quality. The example uses `MLproject` to set up a
Conda environment, define parameter types and defaults, entry point for training, etc.
- `sklearn_logistic_regression` is a simple MLflow example with hooks to log training data to MLflow
tracking server.
- `supply_chain_security` shows how to strengthen the security of ML projects against supply-chain attacks by enforcing hash checks on Python packages.
- `tensorflow` contains end-to-end one run examples from train to predict for TensorFlow 2.8+ It includes usage of MLflow's
`mlflow.tensorflow.autolog()` API, which captures TensorBoard data and logs to MLflow with no code change.
- `docker` demonstrates how to create and run an MLflow project using docker (rather than conda)
to manage project dependencies
- `johnsnowlabs` gives you access to [20.000+ state-of-the-art enterprise NLP models in 200+ languages](https://nlp.johnsnowlabs.com/models) for medical, finance, legal and many more domains.
## Demos
- `demos` folder contains notebooks used during MLflow presentations.
+63
View File
@@ -0,0 +1,63 @@
"""
This is an example for leveraging MLflow's auto tracing capabilities for AutoGen.
For more information about MLflow Tracing, see: https://mlflow.org/docs/latest/llms/tracing/index.html
"""
import os
from typing import Annotated, Literal
from autogen import ConversableAgent
import mlflow
# Turn on auto tracing for AutoGen by calling mlflow.autogen.autolog()
mlflow.autogen.autolog()
config_list = [
{
"model": "gpt-4o-mini",
# Please set your OpenAI API Key to the OPENAI_API_KEY env var before running this example
"api_key": os.environ.get("OPENAI_API_KEY"),
}
]
Operator = Literal["+", "-", "*", "/"]
def calculator(a: int, b: int, operator: Annotated[Operator, "operator"]) -> int:
if operator == "+":
return a + b
elif operator == "-":
return a - b
elif operator == "*":
return a * b
elif operator == "/":
return int(a / b)
else:
raise ValueError("Invalid operator")
# First define the assistant agent that suggests tool calls.
assistant = ConversableAgent(
name="Assistant",
system_message="You are a helpful AI assistant. "
"You can help with simple calculations. "
"Return 'TERMINATE' when the task is done.",
llm_config={"config_list": config_list},
)
# The user proxy agent is used for interacting with the assistant agent
# and executes tool calls.
user_proxy = ConversableAgent(
name="ToolAgent",
llm_config=False,
is_termination_msg=lambda msg: msg.get("content") is not None and "TERMINATE" in msg["content"],
human_input_mode="NEVER",
)
# Register the tool signature with the assistant agent.
assistant.register_for_llm(name="calculator", description="A simple calculator")(calculator)
user_proxy.register_for_execution(name="calculator")(calculator)
response = user_proxy.initiate_chat(assistant, message="What is (44231 + 13312 / (230 - 20)) * 4?")
+76
View File
@@ -0,0 +1,76 @@
import mlflow
mlflow.set_tracking_uri("http://localhost:5000")
mlflow.set_experiment("AGNO Reasoning Finance Team")
mlflow.agno.autolog()
mlflow.anthropic.autolog()
mlflow.openai.autolog()
from agno.agent import Agent
from agno.models.anthropic import Claude
from agno.models.openai import OpenAIChat
from agno.team.team import Team
from agno.tools.duckduckgo import DuckDuckGoTools
from agno.tools.reasoning import ReasoningTools
from agno.tools.yfinance import YFinanceTools
web_agent = Agent(
name="Web Search Agent",
role="Handle web search requests and general research",
model=OpenAIChat(id="gpt-4.1"),
tools=[DuckDuckGoTools()],
instructions="Always include sources",
add_datetime_to_instructions=True,
)
finance_agent = Agent(
name="Finance Agent",
role="Handle financial data requests and market analysis",
model=OpenAIChat(id="gpt-4.1"),
tools=[
YFinanceTools(
stock_price=True,
stock_fundamentals=True,
analyst_recommendations=True,
company_info=True,
)
],
instructions=[
"Use tables to display stock prices, fundamentals (P/E, Market Cap), and recommendations.",
"Clearly state the company name and ticker symbol.",
"Focus on delivering actionable financial insights.",
],
add_datetime_to_instructions=True,
)
reasoning_finance_team = Team(
name="Reasoning Finance Team",
mode="coordinate",
model=Claude(id="claude-sonnet-4-20250514"),
members=[web_agent, finance_agent],
tools=[ReasoningTools(add_instructions=True)],
instructions=[
"Collaborate to provide comprehensive financial and investment insights",
"Consider both fundamental analysis and market sentiment",
"Use tables and charts to display data clearly and professionally",
"Present findings in a structured, easy-to-follow format",
"Only output the final consolidated analysis, not individual agent responses",
],
markdown=True,
show_members_responses=True,
enable_agentic_context=True,
add_datetime_to_instructions=True,
success_criteria="The team has provided a complete financial analysis with data, visualizations, risk assessment, and actionable investment recommendations supported by quantitative analysis and market research.",
)
if __name__ == "__main__":
reasoning_finance_team.print_response(
"""Compare the tech sector giants (AAPL, GOOGL, MSFT) performance:
1. Get financial data for all three companies
2. Analyze recent news affecting the tech sector
3. Calculate comparative metrics and correlations
4. Recommend portfolio allocation weights""",
stream=False,
show_full_reasoning=True,
)
+27
View File
@@ -0,0 +1,27 @@
"""
This is an example for leveraging MLflow's auto tracing capabilities for Anthropic.
For more information about MLflow Tracing, see: https://mlflow.org/docs/latest/llms/tracing/index.html
"""
import os
import anthropic
import mlflow
# Turn on auto tracing for Anthropic by calling mlflow.anthropic.autolog()
mlflow.anthropic.autolog()
# Configure your API key.
client = anthropic.Anthropic(api_key=os.environ["ANTHROPIC_API_KEY"])
# Use the create method to create new message.
message = client.messages.create(
model="claude-3-5-sonnet-20241022",
max_tokens=1024,
messages=[
{"role": "user", "content": "Hello, Claude"},
],
)
print(message.content)
+23
View File
@@ -0,0 +1,23 @@
# Basic authentication example
This example demonstrates the authentication and authorization feature of MLflow.
To run this example,
1. Start the tracking server
```shell
mlflow server --app-name=basic-auth
```
2. Go to `http://localhost:5000/signup` and register two users:
- `(user_a, password_a)`
- `(user_b, password_b)`
3. Run the script
```shell
python auth.py
```
Expected output:
```
2023/05/02 14:03:58 INFO mlflow.tracking.fluent: Experiment with name 'experiment_a' does not exist. Creating a new experiment.
{}
API request to endpoint /api/2.0/mlflow/runs/create failed with error code 403 != 200. Response body: 'Permission denied'
```
+66
View File
@@ -0,0 +1,66 @@
import os
import uuid
import mlflow.server
class User:
MLFLOW_TRACKING_USERNAME = "MLFLOW_TRACKING_USERNAME"
MLFLOW_TRACKING_PASSWORD = "MLFLOW_TRACKING_PASSWORD"
def __init__(self, username, password) -> None:
self.username = username
self.password = password
self.env = {}
def _record_env_var(self, key):
if key := os.environ.get(key):
self.env[key] = key
def _restore_env_var(self, key):
if value := self.env.get(key):
os.environ[key] = value
else:
del os.environ[key]
def __enter__(self):
self._record_env_var(User.MLFLOW_TRACKING_USERNAME)
self._record_env_var(User.MLFLOW_TRACKING_PASSWORD)
os.environ[User.MLFLOW_TRACKING_USERNAME] = self.username
os.environ[User.MLFLOW_TRACKING_PASSWORD] = self.password
return self
def __exit__(self, *_exc):
self._restore_env_var(User.MLFLOW_TRACKING_USERNAME)
self._restore_env_var(User.MLFLOW_TRACKING_PASSWORD)
self.env.clear()
tracking_uri = "http://localhost:5000"
mlflow.set_tracking_uri(tracking_uri)
client = mlflow.server.get_app_client("basic-auth", tracking_uri)
A = User("user_a", "password_a")
B = User("user_b", "password_b")
with A:
exp_a = mlflow.set_experiment(uuid.uuid4().hex)
with mlflow.start_run():
mlflow.log_metric("a", 1)
with B:
mlflow.set_experiment(exp_a.name)
try:
with mlflow.start_run(): # not allowed
mlflow.log_metric("b", 2)
except Exception as e:
print(str(e))
# Grant B permission to edit A's experiment
with A:
client.create_experiment_permission(str(exp_a.experiment_id), B.username, "EDIT")
# B can edit now, should be able to log a metric
with B:
mlflow.set_experiment(exp_a.name)
with mlflow.start_run():
mlflow.log_metric("b", 2)
+38
View File
@@ -0,0 +1,38 @@
# Based on the official regression example:
# https://catboost.ai/docs/concepts/python-usages-examples.html#regression
import numpy as np
from catboost import CatBoostRegressor
import mlflow
from mlflow.models import infer_signature
# Initialize data
train_data = np.array([[1, 4, 5, 6], [4, 5, 6, 7], [30, 40, 50, 60]])
train_labels = np.array([10, 20, 30])
eval_data = np.array([[2, 4, 6, 8], [1, 4, 50, 60]])
# Initialize CatBoostRegressor
params = {
"iterations": 2,
"learning_rate": 1,
"depth": 2,
"allow_writing_files": False,
}
model = CatBoostRegressor(**params)
# Fit model
model.fit(train_data, train_labels)
# Log parameters and fitted model
with mlflow.start_run() as run:
signature = infer_signature(eval_data, model.predict(eval_data))
mlflow.log_params(params)
model_info = mlflow.catboost.log_model(model, name="model", signature=signature)
# Load model
loaded_model = mlflow.catboost.load_model(model_info.model_uri)
# Get predictions
preds = loaded_model.predict(eval_data)
print("predictions:", preds)
+143
View File
@@ -0,0 +1,143 @@
"""
This is an example for leveraging MLflow's auto tracing capabilities for CrewAI.
Most codes are from https://github.com/crewAIInc/crewAI-examples/tree/main/trip_planner.
For more information about MLflow Tracing, see: https://mlflow.org/docs/latest/llms/tracing/index.html
Note that the following example works with crewai>=0.83.0.
"""
from textwrap import dedent
from crewai import Agent, Crew, Task
from crewai.knowledge.source.string_knowledge_source import StringKnowledgeSource
from crewai_tools import SerperDevTool, WebsiteSearchTool
import mlflow
mlflow.set_experiment("CrewAI")
# Turn on auto tracing by calling mlflow.crewai.autolog()
mlflow.crewai.autolog()
content = "Users name is John. He is 30 years old and lives in San Francisco."
string_source = StringKnowledgeSource(content=content, metadata={"preference": "personal"})
search_tool = SerperDevTool()
web_rag_tool = WebsiteSearchTool()
class TripAgents:
def city_selection_agent(self):
return Agent(
role="City Selection Expert",
goal="Select the best city based on weather, season, and prices",
backstory="An expert in analyzing travel data to pick ideal destinations",
tools=[search_tool, web_rag_tool],
verbose=True,
)
def local_expert(self):
return Agent(
role="Local Expert at this city",
goal="Provide the BEST insights about the selected city",
backstory="""A knowledgeable local guide with extensive information
about the city, it's attractions and customs""",
tools=[search_tool, web_rag_tool],
verbose=True,
)
class TripTasks:
def identify_task(self, agent, origin, cities, interests, range):
return Task(
description=dedent(f"""
Analyze and select the best city for the trip based
on specific criteria such as weather patterns, seasonal
events, and travel costs. This task involves comparing
multiple cities, considering factors like current weather
conditions, upcoming cultural or seasonal events, and
overall travel expenses.
Your final answer must be a detailed
report on the chosen city, and everything you found out
about it, including the actual flight costs, weather
forecast and attractions.
Traveling from: {origin}
City Options: {cities}
Trip Date: {range}
Traveler Interests: {interests}
"""),
agent=agent,
expected_output="Detailed report on the chosen city including flight costs, weather forecast, and attractions",
)
def gather_task(self, agent, origin, interests, range):
return Task(
description=dedent(f"""
As a local expert on this city you must compile an
in-depth guide for someone traveling there and wanting
to have THE BEST trip ever!
Gather information about key attractions, local customs,
special events, and daily activity recommendations.
Find the best spots to go to, the kind of place only a
local would know.
This guide should provide a thorough overview of what
the city has to offer, including hidden gems, cultural
hotspots, must-visit landmarks, weather forecasts, and
high level costs.
The final answer must be a comprehensive city guide,
rich in cultural insights and practical tips,
tailored to enhance the travel experience.
Trip Date: {range}
Traveling from: {origin}
Traveler Interests: {interests}
"""),
agent=agent,
expected_output="Comprehensive city guide including hidden gems, cultural hotspots, and practical travel tips",
)
class TripCrew:
def __init__(self, origin, cities, date_range, interests):
self.cities = cities
self.origin = origin
self.interests = interests
self.date_range = date_range
def run(self):
agents = TripAgents()
tasks = TripTasks()
city_selector_agent = agents.city_selection_agent()
local_expert_agent = agents.local_expert()
identify_task = tasks.identify_task(
city_selector_agent, self.origin, self.cities, self.interests, self.date_range
)
gather_task = tasks.gather_task(
local_expert_agent, self.origin, self.interests, self.date_range
)
crew = Crew(
agents=[city_selector_agent, local_expert_agent],
tasks=[identify_task, gather_task],
verbose=True,
memory=True,
knowledge={
"collection_name": "crewai_example",
"sources": [string_source],
"metadata": {"preference": "personal"},
},
)
result = crew.kickoff()
return result
trip_crew = TripCrew("California", "Tokyo", "Dec 12 - Dec 20", "sports")
result = trip_crew.run()
print("\n\n########################")
print("## Here is you Trip Plan")
print("########################\n")
print(result)
+56
View File
@@ -0,0 +1,56 @@
"""
python examples/databricks/dbconnect.py --cluster-id <cluster-id>
"""
import argparse
from databricks.connect import DatabricksSession
from databricks.sdk import WorkspaceClient
from pyspark.sql.types import DoubleType
from sklearn import datasets
from sklearn.neighbors import KNeighborsClassifier
import mlflow
from mlflow.models import infer_signature
def parse_args():
parser = argparse.ArgumentParser()
parser.add_argument("--cluster-id", required=True)
return parser.parse_args()
def main() -> None:
args = parse_args()
wc = WorkspaceClient()
# Train a model
X, y = datasets.load_iris(as_frame=True, return_X_y=True)
model = KNeighborsClassifier().fit(X, y)
predictions = model.predict(X)
signature = infer_signature(X, predictions)
# Log the model
mlflow.set_tracking_uri("databricks")
mlflow.set_experiment(f"/Users/{wc.current_user.me().user_name}/dbconnect")
with mlflow.start_run():
model_info = mlflow.sklearn.log_model(model, name="model", signature=signature)
spark = DatabricksSession.builder.remote(
host=wc.config.host,
token=wc.config.token,
cluster_id=args.cluster_id,
).getOrCreate()
sdf = spark.createDataFrame(X.head(5))
pyfunc_udf = mlflow.pyfunc.spark_udf(
spark,
model_info.model_uri,
env_manager="local",
result_type=DoubleType(),
)
preds = sdf.select(pyfunc_udf(*X.columns).alias("preds"))
preds.show()
if __name__ == "__main__":
main()
+56
View File
@@ -0,0 +1,56 @@
"""
Logs MLflow runs in Databricks from an external host.
How to run:
$ python examples/databricks/log_runs.py --host <host> --token <token> --user <user> [--experiment-id 123]
See also:
https://docs.databricks.com/dev-tools/api/latest/authentication.html#generate-a-personal-access-token
"""
import argparse
import os
import uuid
from sklearn import datasets, svm
from sklearn.model_selection import GridSearchCV, ParameterGrid
import mlflow
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--host", help="Databricks workspace URL")
parser.add_argument("--token", help="Databricks personal access token")
parser.add_argument("--user", help="Databricks username")
parser.add_argument(
"--experiment-id",
default=None,
help="ID of the experiment to log runs in. If unspecified, a new experiment will be created.",
)
args = parser.parse_args()
os.environ["DATABRICKS_HOST"] = args.host
os.environ["DATABRICKS_TOKEN"] = args.token
mlflow.set_tracking_uri("databricks")
if args.experiment_id:
experiment = mlflow.set_experiment(experiment_id=args.experiment_id)
else:
experiment = mlflow.set_experiment(f"/Users/{args.user}/{uuid.uuid4().hex}")
print(f"Logging runs in {args.host}#/mlflow/experiments/{experiment.experiment_id}")
mlflow.sklearn.autolog(max_tuning_runs=None)
iris = datasets.load_iris()
parameters = {"kernel": ("linear", "rbf"), "C": [1, 5, 10]}
clf = GridSearchCV(svm.SVC(), parameters)
clf.fit(iris.data, iris.target)
# Log unnested runs
for params in ParameterGrid(parameters):
clf = svm.SVC(**params)
clf.fit(iris.data, iris.target)
if __name__ == "__main__":
main()
+148
View File
@@ -0,0 +1,148 @@
"""
Benchmark for multi-part upload and download of artifacts.
"""
import hashlib
import json
import os
import pathlib
import tempfile
from concurrent.futures import ThreadPoolExecutor, as_completed
import pandas as pd
import psutil
from tqdm.auto import tqdm
import mlflow
from mlflow.environment_variables import (
MLFLOW_ENABLE_MULTIPART_DOWNLOAD,
MLFLOW_ENABLE_MULTIPART_UPLOAD,
)
from mlflow.utils.time import Timer
GiB = 1024**3
def show_system_info():
svmem = psutil.virtual_memory()
info = json.dumps(
{
"MLflow version": mlflow.__version__,
"MPU enabled": MLFLOW_ENABLE_MULTIPART_DOWNLOAD.get(),
"MPD enabled": MLFLOW_ENABLE_MULTIPART_UPLOAD.get(),
"CPU count": psutil.cpu_count(),
"Memory usage (total) [GiB]": svmem.total // GiB,
"Memory used [GiB]": svmem.used // GiB,
"Memory available [GiB]": svmem.available // GiB,
},
indent=2,
)
max_len = max(map(len, info.splitlines()))
print("=" * max_len)
print(info)
print("=" * max_len)
def md5_checksum(path):
file_hash = hashlib.sha256()
with open(path, "rb") as f:
while chunk := f.read(1024**2):
file_hash.update(chunk)
return file_hash.hexdigest()
def assert_checksum_equal(path1, path2):
assert md5_checksum(path1) == md5_checksum(path2), f"Checksum mismatch for {path1} and {path2}"
def yield_random_bytes(num_bytes):
while num_bytes > 0:
chunk_size = min(num_bytes, 1024**2)
yield os.urandom(chunk_size)
num_bytes -= chunk_size
def generate_random_file(path, num_bytes):
with open(path, "wb") as f:
for chunk in yield_random_bytes(num_bytes):
f.write(chunk)
def upload_and_download(file_size, num_files):
with tempfile.TemporaryDirectory() as tmpdir:
tmpdir = pathlib.Path(tmpdir)
# Prepare files
src_dir = tmpdir / "src"
src_dir.mkdir()
files = {}
with ThreadPoolExecutor() as pool:
futures = []
for i in range(num_files):
f = src_dir / str(i)
futures.append(pool.submit(generate_random_file, f, file_size))
files[f.name] = f
for fut in tqdm(
as_completed(futures),
total=len(futures),
desc="Generating files",
colour="#FFA500",
):
fut.result()
# Upload
with mlflow.start_run() as run:
with Timer() as t_upload:
mlflow.log_artifacts(str(src_dir))
# Download
dst_dir = tmpdir / "dst"
dst_dir.mkdir()
with Timer() as t_download:
mlflow.artifacts.download_artifacts(
artifact_uri=f"{run.info.artifact_uri}/", dst_path=dst_dir
)
# Verify checksums
with ThreadPoolExecutor() as pool:
futures = []
for f in dst_dir.rglob("*"):
if f.is_dir():
continue
futures.append(pool.submit(assert_checksum_equal, f, files[f.name]))
for fut in tqdm(
as_completed(futures),
total=len(futures),
desc="Verifying checksums",
colour="#FFA500",
):
fut.result()
return t_upload.elapsed, t_download.elapsed
def main():
# Uncomment the following lines if you're running this script outside of Databricks
# using a personal access token:
# mlflow.set_tracking_uri("databricks")
# mlflow.set_experiment("/Users/<username>/benchmark")
FILE_SIZE = 1 * GiB
NUM_FILES = 2
NUM_ATTEMPTS = 3
show_system_info()
stats = []
for i in range(NUM_ATTEMPTS):
print(f"Attempt {i + 1} / {NUM_ATTEMPTS}")
stats.append(upload_and_download(FILE_SIZE, NUM_FILES))
df = pd.DataFrame(stats, columns=["upload [s]", "download [s]"])
# show mean, min, max in markdown table
print(df.aggregate(["count", "mean", "min", "max"]).to_markdown())
if __name__ == "__main__":
main()
+9
View File
@@ -0,0 +1,9 @@
# MLflow 3 Examples
## Pre-requisites
Before running the examples, run the following command to install mlflow 3.0:
```sh
pip install git+https://github.com/mlflow/mlflow.git@mlflow-3
```
+125
View File
@@ -0,0 +1,125 @@
# # MLflow 3 Deep Learning Example
# In this example, we will first run a model training job, which is tracked as an MLflow Run.
# Every 10 epochs, we will store model checkpoints, which are tracked as MLflow Logged Models.
# We will then select the best checkpoint for production deployment.
import pandas as pd
import torch
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from torch import nn
import mlflow
import mlflow.pytorch
from mlflow.entities import Dataset
# Helper function to prepare data
def prepare_data(df):
X = torch.tensor(df.iloc[:, :-1].values, dtype=torch.float32)
y = torch.tensor(df.iloc[:, -1].values, dtype=torch.long)
return X, y
# Helper function to compute accuracy
def compute_accuracy(model, X, y):
with torch.no_grad():
outputs = model(X)
_, predicted = torch.max(outputs, 1)
accuracy = (predicted == y).sum().item() / y.size(0)
return accuracy
# Define a basic PyTorch classifier
class IrisClassifier(nn.Module):
def __init__(self, input_size, hidden_size, output_size):
super().__init__()
self.fc1 = nn.Linear(input_size, hidden_size)
self.relu = nn.ReLU()
self.fc2 = nn.Linear(hidden_size, output_size)
def forward(self, x):
x = self.fc1(x)
x = self.relu(x)
x = self.fc2(x)
return x
# Load Iris dataset and prepare the DataFrame
iris = load_iris()
iris_df = pd.DataFrame(data=iris.data, columns=iris.feature_names)
iris_df["target"] = iris.target
# Split into training and testing datasets
train_df, test_df = train_test_split(iris_df, test_size=0.2, random_state=42)
# Prepare training data
train_dataset = mlflow.data.from_pandas(train_df, name="train")
X_train, y_train = prepare_data(train_dataset.df)
# Define the PyTorch model and move it to the device
input_size = X_train.shape[1]
hidden_size = 16
output_size = len(iris.target_names)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
scripted_model = IrisClassifier(input_size, hidden_size, output_size).to(device)
scripted_model = torch.jit.script(scripted_model)
# Start a run to represent the training job
with mlflow.start_run():
# Load the training dataset with MLflow. We will link training metrics to this dataset.
train_dataset: Dataset = mlflow.data.from_pandas(train_df, name="train")
X_train, y_train = prepare_data(train_dataset.df)
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(scripted_model.parameters(), lr=0.01)
for epoch in range(51):
X_train = X_train.to(device)
y_train = y_train.to(device)
out = scripted_model(X_train)
loss = criterion(out, y_train)
optimizer.zero_grad()
loss.backward()
optimizer.step()
# Log a checkpoint with metrics every 10 epochs
if epoch % 10 == 0:
# Each newly created LoggedModel checkpoint is linked with its
# name, params, and step
model_info = mlflow.pytorch.log_model(
pytorch_model=scripted_model,
name=f"torch-iris-{epoch}",
params={
"n_layers": 3,
"activation": "ReLU",
"criterion": "CrossEntropyLoss",
"optimizer": "Adam",
},
step=epoch,
input_example=X_train.numpy(),
)
# Log metric on training dataset at step and link to LoggedModel
mlflow.log_metric(
key="accuracy",
value=compute_accuracy(scripted_model, X_train, y_train),
step=epoch,
model_id=model_info.model_id,
dataset=train_dataset,
)
# This example produced one MLflow Run (training_run) and 6 MLflow Logged Models,
# one for each checkpoint (at steps 0, 10, …, 50). Using MLflow's UI or search API,
# we can get the checkpoints and rank them by their accuracy.
ranked_checkpoints = mlflow.search_logged_models(
output_format="list", order_by=[{"field_name": "metrics.accuracy", "ascending": False}]
)
best_checkpoint: mlflow.entities.LoggedModel = ranked_checkpoints[0]
print(best_checkpoint.metrics[0])
print(best_checkpoint)
worst_checkpoint: mlflow.entities.LoggedModel = ranked_checkpoints[-1]
print(worst_checkpoint.metrics)
# Once the best checkpoint is selected, that model can be registered to the model registry.
mlflow.register_model(f"models:/{best_checkpoint.model_id}", name="my_dl_model")
+46
View File
@@ -0,0 +1,46 @@
# MLflow 3 GenAI Example
# In this example, we will create an agent and then evaluate its performance. First, we will define the agent and log it to MLflow.
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI
import mlflow
# Define the chain
chat_model = ChatOpenAI(name="gpt-4o")
prompt = ChatPromptTemplate.from_messages([
("system", "You are a chatbot that can answer questions about Databricks."),
("user", "{messages}"),
])
chain = prompt | chat_model
# Log the chain with MLflow, specifying its parameters
# As a new feature, the LoggedModel entity is linked to its name and params
logged_model = mlflow.langchain.log_model(
lc_model=chain,
name="basic_chain",
params={"temperature": 0.1, "max_tokens": 2000, "prompt_template": str(prompt)},
model_type="agent",
input_example={"messages": "What is MLflow?"},
)
# Inspect the LoggedModel and its properties
print(logged_model.model_id, logged_model.params)
# m-123802d4ba324f4d8baa456eb8b5c061, {'max_tokens': '2000', 'prompt_template': "input_variables=['messages'] messages=[SystemMessagePromptTemplate(...), HumanMessagePromptTemplate(...)]", 'temperature': '0.1'}
# Then, we will interactively query the chain in a notebook to make sure that it's viable enough for further testing. These traces can be viewed in UI, under the Traces tab of the model details page.
# Enable autologging so that interactive traces from the chain are automatically linked to its LoggedModel
mlflow.langchain.autolog()
loaded_chain = mlflow.langchain.load_model(f"models:/{logged_model.model_id}")
chain_inputs = [
{"messages": "What is MLflow?"},
{"messages": "What is Unity Catalog?"},
{"messages": "What are user-defined functions (UDFs)?"},
]
for chain_input in chain_inputs:
loaded_chain.invoke(chain_input)
# Print out the traces linked to the LoggedModel
print(mlflow.search_traces(model_id=logged_model.model_id))
+113
View File
@@ -0,0 +1,113 @@
# MLflow 3 Traditional ML Example
# In this example, we will first run a model training job, which is tracked as
# an MLflow Run, to produce a trained model, which is tracked as an MLflow Logged Model.
import pandas as pd
from sklearn.datasets import load_iris
from sklearn.linear_model import ElasticNet
from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score
from sklearn.model_selection import train_test_split
import mlflow
import mlflow.sklearn
from mlflow.entities import Dataset
# Helper function to compute metrics
def compute_metrics(actual, predicted):
rmse = mean_squared_error(actual, predicted)
mae = mean_absolute_error(actual, predicted)
r2 = r2_score(actual, predicted)
return rmse, mae, r2
# Load Iris dataset and prepare the DataFrame
iris = load_iris()
iris_df = pd.DataFrame(data=iris.data, columns=iris.feature_names)
iris_df["quality"] = (iris.target == 2).astype(int) # Create a binary target for simplicity
# Split into training and testing datasets
train_df, test_df = train_test_split(iris_df, test_size=0.2, random_state=42)
# Start a run to represent the training job
with mlflow.start_run() as training_run:
# Load the training dataset with MLflow. We will link training metrics to this dataset.
train_dataset: Dataset = mlflow.data.from_pandas(train_df, name="train")
train_x = train_dataset.df.drop(["quality"], axis=1)
train_y = train_dataset.df[["quality"]]
# Fit a model to the training dataset
lr = ElasticNet(alpha=0.5, l1_ratio=0.5, random_state=42)
lr.fit(train_x, train_y)
# Log the model, specifying its ElasticNet parameters (alpha, l1_ratio)
# As a new feature, the LoggedModel entity is linked to its name and params
logged_model = mlflow.sklearn.log_model(
sk_model=lr,
name="elasticnet",
params={
"alpha": 0.5,
"l1_ratio": 0.5,
},
input_example=train_x,
)
# Inspect the LoggedModel and its properties
print(logged_model.model_id, logged_model.params)
# m-fa4e1bca8cb64971bce2322a8fd427d3, {'alpha': '0.5', 'l1_ratio': '0.5'}
# Evaluate the model on the training dataset and log metrics
# These metrics are now linked to the LoggedModel entity
predictions = lr.predict(train_x)
(rmse, mae, r2) = compute_metrics(train_y, predictions)
mlflow.log_metrics(
metrics={
"rmse": rmse,
"r2": r2,
"mae": mae,
},
model_id=logged_model.model_id,
dataset=train_dataset,
)
# Inspect the LoggedModel, now with metrics
logged_model = mlflow.get_logged_model(logged_model.model_id)
print(logged_model.model_id, logged_model.metrics)
# m-fa4e1bca8cb64971bce2322a8fd427d3, [<Metric: dataset_name='train', key='rmse', model_id='m-fa4e1bca8cb64971bce2322a8fd427d3, value=0.7538635773139717, ...>, ...]
# Some time later, when we get a new evaluation dataset based on the latest production data,
# we will run a new model evaluation job, which is tracked as a new MLflow Run,
# to measure the performance of the model on this new dataset.
# This example will produced two MLflow Runs (training_run and evaluation_run) and
# one MLflow Logged Model (elasticnet). From the resulting Logged Model,
# we can see all of the parameters and metadata. We can also see all of the metrics linked
# from the training and evaluation runs.
# Start a run to represent the test dataset evaluation job
with mlflow.start_run() as evaluation_run:
# Load the test dataset with MLflow. We will link test metrics to this dataset.
test_dataset: mlflow.entities.Dataset = mlflow.data.from_pandas(test_df, name="test")
test_x = test_dataset.df.drop(["quality"], axis=1)
test_y = test_dataset.df[["quality"]]
# Load the model
model = mlflow.sklearn.load_model(f"models:/{logged_model.model_id}")
# Evaluate the model on the training dataset and log metrics, linking to model
predictions = model.predict(test_x)
(rmse, mae, r2) = compute_metrics(test_y, predictions)
mlflow.log_metrics(
metrics={
"rmse": rmse,
"r2": r2,
"mae": mae,
},
dataset=test_dataset,
model_id=logged_model.model_id,
)
print(mlflow.get_logged_model(logged_model.model_id).to_dictionary())
# Now register the model.
mlflow.register_model(logged_model.model_uri, name="my_ml_model")
File diff suppressed because one or more lines are too long
@@ -0,0 +1,450 @@
{
"cells": [
{
"cell_type": "markdown",
"id": "66132127-faa2-43d7-a23d-94035f0dc6a4",
"metadata": {},
"source": [
"# PythonModel with type hints Quickstart\n",
"This notebook will demonstrates how to use type hints for data validation in MLflow PythonModel.\n",
"\n",
"## Prerequisite\n",
"\n",
"Install required packages: `pip install mlflow==2.20.0 openai==1.65.4`\n",
"\n",
"Set your OpenAI API key with `os.environ[\"OPENAI_API_KEY\"]=\"<YOUR_KEY>\"`"
]
},
{
"cell_type": "markdown",
"id": "3ce58023-e1ab-4bf9-8325-0e706e327c54",
"metadata": {},
"source": [
"## Create a model"
]
},
{
"cell_type": "code",
"execution_count": 1,
"id": "30503174-848c-4017-85f6-87f1b90fce70",
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"ChatCompletion(id='chatcmpl-B83LveI6Wc2RHEdbMwNGkFh8ZoehY', choices=[Choice(finish_reason='stop', index=0, logprobs=None, message=ChatCompletionMessage(content='MLflow is an open-source platform designed to manage the machine learning (ML) lifecycle, which includes components such as experimentation, reproducibility, and deployment. It provides a suite of tools to help data scientists and machine learning practitioners track experiments, package and share their code, and deploy models. Here are the key components of MLflow:\\n\\n1. **MLflow Tracking**: This component allows users to log and query experiments. You can track metrics, parameters, and artifacts (such as model files) associated with different runs of your machine learning models.\\n\\n2. **MLflow Projects**: This functionality enables users to package their data science code in a reusable and reproducible way. Projects are defined with a standard format that specifies their dependencies and how to run them.\\n\\n3. **MLflow Models**: This part of MLflow provides a way to manage and deploy machine learning models in various formats. It supports multiple model types, enabling users to deploy models in different environments, such as REST APIs, cloud services, or on-premises servers.\\n\\n4. **MLflow Registry**: This is a centralized model store that manages the lifecycle of machine learning models, including versioning, stage transitions (like staging, production, archived), and annotations. It allows teams to track model changes and collaborate more effectively.\\n\\nMLflow is designed to be flexible and integrates well with existing machine learning frameworks like TensorFlow, PyTorch, Scikit-Learn, and many others. This flexibility makes it widely adopted in production ML workflows and among research communities. \\n\\nOverall, MLflow aims to streamline the process of developing, tracking, and deploying machine learning models, reducing friction and enhancing collaboration among data science teams.', refusal=None, role='assistant', audio=None, function_call=None, tool_calls=None))], created=1741259211, model='gpt-4o-mini-2024-07-18', object='chat.completion', service_tier='default', system_fingerprint='fp_06737a9306', usage=CompletionUsage(completion_tokens=339, prompt_tokens=13, total_tokens=352, completion_tokens_details=CompletionTokensDetails(accepted_prediction_tokens=0, audio_tokens=0, reasoning_tokens=0, rejected_prediction_tokens=0), prompt_tokens_details=PromptTokensDetails(audio_tokens=0, cached_tokens=0)))"
]
},
"execution_count": 1,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"import openai\n",
"import pydantic\n",
"\n",
"import mlflow\n",
"\n",
"# Use OpenAI model\n",
"messages = [{\"role\": \"user\", \"content\": \"What is the MLflow?\"}]\n",
"response = openai.chat.completions.create(model=\"gpt-4o-mini\", messages=messages)\n",
"response"
]
},
{
"cell_type": "markdown",
"id": "6f1522dd-089e-4908-be5d-08de29deb683",
"metadata": {},
"source": [
"## Use MLflow PythonModel with type hints"
]
},
{
"cell_type": "code",
"execution_count": 2,
"id": "7a960b6b-3dc7-47e3-ac4e-a50ef296cae6",
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"'DSPy is a Python library designed for creating and managing decision systems, particularly in the context of data-driven applications. It provides tools for building models that can make decisions based on inputs from various data sources. The library aims to simplify the process of developing and deploying decision logic, which is often a complex task in machine learning and artificial intelligence projects.\\n\\nKey features of DSPy may include:\\n\\n1. **Declarative Syntax**: Allowing users to express decision logic in a clear and concise manner.\\n2. **Integration with Data Sources**: Facilitating easy integration with various data workflows, making it simpler to utilize datasets for decision-making.\\n3. **Evaluation and Testing**: Providing tools for evaluating and testing decision-making models, ensuring their accuracy and reliability.\\n\\nBy leveraging DSPy, data scientists and developers can focus on building effective decision systems without getting bogged down by the complexities usually associated with programming these systems from scratch.\\n\\nFor the latest updates and more specific functionalities, it's a good idea to refer to the official DSPy documentation or repository, as libraries are frequently updated and improved.'"
]
},
"execution_count": 2,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"# define your input schema\n",
"class Message(pydantic.BaseModel):\n",
" role: str\n",
" content: str\n",
"\n",
"\n",
"# inherit mlflow PythonModel\n",
"class MyModel(mlflow.pyfunc.PythonModel):\n",
" # add type hint to model_input\n",
" def predict(self, model_input: list[Message]) -> str:\n",
" response = openai.chat.completions.create(model=\"gpt-4o-mini\", messages=model_input)\n",
" return response.choices[0].message.content\n",
"\n",
"\n",
"model = MyModel()\n",
"model.predict([{\"role\": \"user\", \"content\": \"What is DSPy?\"}])"
]
},
{
"cell_type": "code",
"execution_count": 3,
"id": "15d8e83c-3afa-4b75-80bf-1a8efa51a415",
"metadata": {
"scrolled": true
},
"outputs": [
{
"ename": "MlflowException",
"evalue": "Failed to validate data against type hint `list[Message]`, invalid elements: [('What is DSPy?', \"Expecting example to be a dictionary or pydantic model instance for Pydantic type hint, got <class 'str'>\")]",
"output_type": "error",
"traceback": [
"\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
"\u001b[0;31mMlflowException\u001b[0m Traceback (most recent call last)",
"Cell \u001b[0;32mIn[3], line 2\u001b[0m\n\u001b[1;32m 1\u001b[0m \u001b[38;5;66;03m# An incorrect input will trigger validation error\u001b[39;00m\n\u001b[0;32m----> 2\u001b[0m \u001b[43mmodel\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mpredict\u001b[49m\u001b[43m(\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mWhat is DSPy?\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m]\u001b[49m\u001b[43m)\u001b[49m\n",
"File \u001b[0;32m~/Documents/repos/mlflow/mlflow/pyfunc/utils/data_validation.py:68\u001b[0m, in \u001b[0;36m_wrap_predict_with_pyfunc.<locals>.wrapper\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m 66\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mException\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m e:\n\u001b[1;32m 67\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(e, MlflowException):\n\u001b[0;32m---> 68\u001b[0m \u001b[38;5;28;01mraise\u001b[39;00m e\n\u001b[1;32m 69\u001b[0m \u001b[38;5;28;01mraise\u001b[39;00m MlflowException(\n\u001b[1;32m 70\u001b[0m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mFailed to validate the input data against the type hint \u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m 71\u001b[0m \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m`\u001b[39m\u001b[38;5;132;01m{\u001b[39;00mfunc_info\u001b[38;5;241m.\u001b[39minput_type_hint\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m`. Error: \u001b[39m\u001b[38;5;132;01m{\u001b[39;00me\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m 72\u001b[0m )\n\u001b[1;32m 73\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m func(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n",
"File \u001b[0;32m~/Documents/repos/mlflow/mlflow/pyfunc/utils/data_validation.py:59\u001b[0m, in \u001b[0;36m_wrap_predict_with_pyfunc.<locals>.wrapper\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m 56\u001b[0m \u001b[38;5;129m@wraps\u001b[39m(func)\n\u001b[1;32m 57\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mwrapper\u001b[39m(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs):\n\u001b[1;32m 58\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m---> 59\u001b[0m args, kwargs \u001b[38;5;241m=\u001b[39m \u001b[43m_validate_model_input\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m 60\u001b[0m \u001b[43m \u001b[49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 61\u001b[0m \u001b[43m \u001b[49m\u001b[43mkwargs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 62\u001b[0m \u001b[43m \u001b[49m\u001b[43mmodel_input_index\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 63\u001b[0m \u001b[43m \u001b[49m\u001b[43mfunc_info\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43minput_type_hint\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 64\u001b[0m \u001b[43m \u001b[49m\u001b[43mfunc_info\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43minput_param_name\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 65\u001b[0m \u001b[43m \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 66\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mException\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m e:\n\u001b[1;32m 67\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(e, MlflowException):\n",
"File \u001b[0;32m~/Documents/repos/mlflow/mlflow/pyfunc/utils/data_validation.py:200\u001b[0m, in \u001b[0;36m_validate_model_input\u001b[0;34m(args, kwargs, model_input_index_in_sig, type_hint, model_input_param_name)\u001b[0m\n\u001b[1;32m 198\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m input_pos \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[1;32m 199\u001b[0m data \u001b[38;5;241m=\u001b[39m _convert_data_to_type_hint(model_input, type_hint)\n\u001b[0;32m--> 200\u001b[0m data \u001b[38;5;241m=\u001b[39m \u001b[43m_validate_data_against_type_hint\u001b[49m\u001b[43m(\u001b[49m\u001b[43mdata\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtype_hint\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 201\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m input_pos \u001b[38;5;241m==\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mkwargs\u001b[39m\u001b[38;5;124m\"\u001b[39m:\n\u001b[1;32m 202\u001b[0m kwargs[model_input_param_name] \u001b[38;5;241m=\u001b[39m data\n",
"File \u001b[0;32m~/Documents/repos/mlflow/mlflow/types/type_hints.py:452\u001b[0m, in \u001b[0;36m_validate_data_against_type_hint\u001b[0;34m(data, type_hint)\u001b[0m\n\u001b[1;32m 450\u001b[0m args \u001b[38;5;241m=\u001b[39m get_args(type_hint)\n\u001b[1;32m 451\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m origin_type \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28mlist\u001b[39m:\n\u001b[0;32m--> 452\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43m_validate_list_elements\u001b[49m\u001b[43m(\u001b[49m\u001b[43melement_type\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;241;43m0\u001b[39;49m\u001b[43m]\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdata\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mdata\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 453\u001b[0m \u001b[38;5;28;01melif\u001b[39;00m origin_type \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28mdict\u001b[39m:\n\u001b[1;32m 454\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m _validate_dict_elements(element_type\u001b[38;5;241m=\u001b[39margs[\u001b[38;5;241m1\u001b[39m], data\u001b[38;5;241m=\u001b[39mdata)\n",
"File \u001b[0;32m~/Documents/repos/mlflow/mlflow/types/type_hints.py:528\u001b[0m, in \u001b[0;36m_validate_list_elements\u001b[0;34m(element_type, data)\u001b[0m\n\u001b[1;32m 524\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m invalid_elems:\n\u001b[1;32m 525\u001b[0m invalid_elems_msg \u001b[38;5;241m=\u001b[39m (\n\u001b[1;32m 526\u001b[0m \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;132;01m{\u001b[39;00minvalid_elems[:\u001b[38;5;241m5\u001b[39m]\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m ... (truncated)\u001b[39m\u001b[38;5;124m\"\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mlen\u001b[39m(invalid_elems) \u001b[38;5;241m>\u001b[39m \u001b[38;5;241m5\u001b[39m \u001b[38;5;28;01melse\u001b[39;00m invalid_elems\n\u001b[1;32m 527\u001b[0m )\n\u001b[0;32m--> 528\u001b[0m \u001b[38;5;28;01mraise\u001b[39;00m MlflowException\u001b[38;5;241m.\u001b[39minvalid_parameter_value(\n\u001b[1;32m 529\u001b[0m \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mFailed to validate data against type hint `list[\u001b[39m\u001b[38;5;132;01m{\u001b[39;00m_type_hint_repr(element_type)\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m]`, \u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m 530\u001b[0m \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124minvalid elements: \u001b[39m\u001b[38;5;132;01m{\u001b[39;00minvalid_elems_msg\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m 531\u001b[0m )\n\u001b[1;32m 532\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m result\n",
"\u001b[0;31mMlflowException\u001b[0m: Failed to validate data against type hint `list[Message]`, invalid elements: [('What is DSPy?', \"Expecting example to be a dictionary or pydantic model instance for Pydantic type hint, got <class 'str'>\")]"
]
}
],
"source": [
"# An incorrect input will trigger validation error\n",
"model.predict([\"What is DSPy?\"])"
]
},
{
"cell_type": "markdown",
"id": "9e5b9c51-d840-435b-9923-c4088087020b",
"metadata": {},
"source": [
"## Model logging"
]
},
{
"cell_type": "code",
"execution_count": 4,
"id": "c4bde4ec-992c-432d-aefa-3087e4a59861",
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"2025/03/06 19:07:07 INFO mlflow.models.signature: Inferring model signature from type hints\n",
"2025/03/06 19:07:07 INFO mlflow.models.signature: Running the predict function to generate output based on input example\n"
]
}
],
"source": [
"# log the model\n",
"with mlflow.start_run():\n",
" model_info = mlflow.pyfunc.log_model(name=\"model\", python_model=model, input_example=messages)"
]
},
{
"cell_type": "code",
"execution_count": 5,
"id": "805d1dd8-b63a-4236-9f7e-1ab845c7a39c",
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"inputs: \n",
" [{content: string (required), role: string (required)} (required)]\n",
"outputs: \n",
" [string (required)]\n",
"params: \n",
" None"
]
},
"execution_count": 5,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"model_info.signature"
]
},
{
"cell_type": "code",
"execution_count": 6,
"id": "ca4e0a8a-924d-4b84-88d3-4a1faa2c01a8",
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"'MLflow is an open-source platform designed to manage the machine learning lifecycle. It provides tools for various stages of the machine learning process, including:\\n\\n1. **Experiment Tracking**: MLflow allows you to log and track experiments, enabling you to compare different runs and their performance metrics easily. You can log parameters, metrics, tags, and artifacts related to your models.\\n\\n2. **Projects**: MLflow Projects facilitate packaging and sharing code in a reusable format. This makes it easier to reproduce experiments and share your work with others.\\n\\n3. **Models**: MLflow Models provides a standard format for packaging machine learning models. It supports various flavors of models (e.g., TensorFlow, PyTorch, Scikit-learn) and allows you to deploy them to various environments (like Docker, cloud-based services, or local servers).\\n\\n4. **Registry**: The MLflow Model Registry provides a centralized repository to manage models, including versioning, annotation, and lifecycle management (staging, production, and archived statuses).\\n\\n5. **Integration**: MLflow integrates well with popular machine learning frameworks and libraries, making it a versatile choice for data scientists and machine learning engineers.\\n\\nBy using MLflow, teams can streamline their machine learning workflows, enhance collaboration, and ensure reproducibility in their experiments, leading to more efficient model development and deployment.'"
]
},
"execution_count": 6,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"# load the pyfunc model\n",
"pyfunc_model = mlflow.pyfunc.load_model(model_info.model_uri)\n",
"# the same validation works for pyfunc model predict\n",
"pyfunc_model.predict(messages)"
]
},
{
"cell_type": "markdown",
"id": "b920f0d1-4c89-43fb-bdc9-e69702b059b3",
"metadata": {},
"source": [
"## Verify model before deployment"
]
},
{
"cell_type": "code",
"execution_count": 11,
"id": "c702c097-19aa-4530-a075-fd5c676a0d9e",
"metadata": {},
"outputs": [
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "c41f48f443de4b5c810b46704358b815",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"Downloading artifacts: 0%| | 0/7 [00:00<?, ?it/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"name": "stderr",
"output_type": "stream",
"text": [
"2025/03/06 19:10:16 INFO mlflow.models.flavor_backend_registry: Selected backend for flavor 'python_function'\n",
"2025/03/06 19:10:16 INFO mlflow.utils.virtualenv: Creating a new environment in /var/folders/9g/psrbbvm92t712cy09d7_00d00000gp/T/tmpp18g94f9/envs/virtualenv_envs/mlflow-9d81fff15053e6e06e2edaefcc9e075d6c04a094 with python version 3.9.18 using uv\n",
"Using CPython 3.9.18 interpreter at: \u001b[36m/Users/serena.ruan/miniconda3/envs/mlflow/bin/python3.9\u001b[39m\n",
"Creating virtual environment at: \u001b[36m/var/folders/9g/psrbbvm92t712cy09d7_00d00000gp/T/tmpp18g94f9/envs/virtualenv_envs/mlflow-9d81fff15053e6e06e2edaefcc9e075d6c04a094\u001b[39m\n",
"Activate with: \u001b[32msource /var/folders/9g/psrbbvm92t712cy09d7_00d00000gp/T/tmpp18g94f9/envs/virtualenv_envs/mlflow-9d81fff15053e6e06e2edaefcc9e075d6c04a094/bin/activate\u001b[39m\n",
"2025/03/06 19:10:16 INFO mlflow.utils.virtualenv: Installing dependencies\n",
"\u001b[2mUsing Python 3.9.18 environment at: /var/folders/9g/psrbbvm92t712cy09d7_00d00000gp/T/tmpp18g94f9/envs/virtualenv_envs/mlflow-9d81fff15053e6e06e2edaefcc9e075d6c04a094\u001b[0m\n",
"\u001b[2mResolved \u001b[1m3 packages\u001b[0m \u001b[2min 1ms\u001b[0m\u001b[0m\n",
"\u001b[2mInstalled \u001b[1m3 packages\u001b[0m \u001b[2min 12ms\u001b[0m\u001b[0m\n",
" \u001b[32m+\u001b[39m \u001b[1mpip\u001b[0m\u001b[2m==23.3\u001b[0m\n",
" \u001b[32m+\u001b[39m \u001b[1msetuptools\u001b[0m\u001b[2m==68.0.0\u001b[0m\n",
" \u001b[32m+\u001b[39m \u001b[1mwheel\u001b[0m\u001b[2m==0.41.2\u001b[0m\n",
"\u001b[2mUsing Python 3.9.18 environment at: /var/folders/9g/psrbbvm92t712cy09d7_00d00000gp/T/tmpp18g94f9/envs/virtualenv_envs/mlflow-9d81fff15053e6e06e2edaefcc9e075d6c04a094\u001b[0m\n",
"\u001b[2mResolved \u001b[1m84 packages\u001b[0m \u001b[2min 4.91s\u001b[0m\u001b[0m\n",
"\u001b[2mPrepared \u001b[1m1 package\u001b[0m \u001b[2min 423ms\u001b[0m\u001b[0m\n",
"\u001b[2mInstalled \u001b[1m83 packages\u001b[0m \u001b[2min 849ms\u001b[0m\u001b[0m\n",
" \u001b[32m+\u001b[39m \u001b[1malembic\u001b[0m\u001b[2m==1.15.1\u001b[0m\n",
" \u001b[32m+\u001b[39m \u001b[1mannotated-types\u001b[0m\u001b[2m==0.7.0\u001b[0m\n",
" \u001b[32m+\u001b[39m \u001b[1manyio\u001b[0m\u001b[2m==4.8.0\u001b[0m\n",
" \u001b[32m+\u001b[39m \u001b[1mattrs\u001b[0m\u001b[2m==23.1.0\u001b[0m\n",
" \u001b[32m+\u001b[39m \u001b[1mblinker\u001b[0m\u001b[2m==1.9.0\u001b[0m\n",
" \u001b[32m+\u001b[39m \u001b[1mbrotli\u001b[0m\u001b[2m==1.1.0\u001b[0m\n",
" \u001b[32m+\u001b[39m \u001b[1mcachetools\u001b[0m\u001b[2m==5.5.2\u001b[0m\n",
" \u001b[32m+\u001b[39m \u001b[1mcertifi\u001b[0m\u001b[2m==2025.1.31\u001b[0m\n",
" \u001b[32m+\u001b[39m \u001b[1mcharset-normalizer\u001b[0m\u001b[2m==3.4.1\u001b[0m\n",
" \u001b[32m+\u001b[39m \u001b[1mclick\u001b[0m\u001b[2m==8.1.7\u001b[0m\n",
" \u001b[32m+\u001b[39m \u001b[1mcloudpickle\u001b[0m\u001b[2m==3.0.0\u001b[0m\n",
" \u001b[32m+\u001b[39m \u001b[1mcontourpy\u001b[0m\u001b[2m==1.3.0\u001b[0m\n",
" \u001b[32m+\u001b[39m \u001b[1mcycler\u001b[0m\u001b[2m==0.12.1\u001b[0m\n",
" \u001b[32m+\u001b[39m \u001b[1mcython\u001b[0m\u001b[2m==3.0.5\u001b[0m\n",
" \u001b[32m+\u001b[39m \u001b[1mdatabricks-sdk\u001b[0m\u001b[2m==0.44.1\u001b[0m\n",
" \u001b[32m+\u001b[39m \u001b[1mdeprecated\u001b[0m\u001b[2m==1.2.18\u001b[0m\n",
" \u001b[32m+\u001b[39m \u001b[1mdistro\u001b[0m\u001b[2m==1.9.0\u001b[0m\n",
" \u001b[32m+\u001b[39m \u001b[1mdocker\u001b[0m\u001b[2m==7.1.0\u001b[0m\n",
" \u001b[32m+\u001b[39m \u001b[1mexceptiongroup\u001b[0m\u001b[2m==1.2.2\u001b[0m\n",
" \u001b[32m+\u001b[39m \u001b[1mflask\u001b[0m\u001b[2m==3.1.0\u001b[0m\n",
" \u001b[32m+\u001b[39m \u001b[1mfonttools\u001b[0m\u001b[2m==4.56.0\u001b[0m\n",
" \u001b[32m+\u001b[39m \u001b[1mgitdb\u001b[0m\u001b[2m==4.0.12\u001b[0m\n",
" \u001b[32m+\u001b[39m \u001b[1mgitpython\u001b[0m\u001b[2m==3.1.44\u001b[0m\n",
" \u001b[32m+\u001b[39m \u001b[1mgoogle-auth\u001b[0m\u001b[2m==2.38.0\u001b[0m\n",
" \u001b[32m+\u001b[39m \u001b[1mgraphene\u001b[0m\u001b[2m==3.4.3\u001b[0m\n",
" \u001b[32m+\u001b[39m \u001b[1mgraphql-core\u001b[0m\u001b[2m==3.2.6\u001b[0m\n",
" \u001b[32m+\u001b[39m \u001b[1mgraphql-relay\u001b[0m\u001b[2m==3.2.0\u001b[0m\n",
" \u001b[32m+\u001b[39m \u001b[1mgunicorn\u001b[0m\u001b[2m==23.0.0\u001b[0m\n",
" \u001b[32m+\u001b[39m \u001b[1mh11\u001b[0m\u001b[2m==0.14.0\u001b[0m\n",
" \u001b[32m+\u001b[39m \u001b[1mh2\u001b[0m\u001b[2m==4.1.0\u001b[0m\n",
" \u001b[32m+\u001b[39m \u001b[1mhpack\u001b[0m\u001b[2m==4.1.0\u001b[0m\n",
" \u001b[32m+\u001b[39m \u001b[1mhttpcore\u001b[0m\u001b[2m==1.0.7\u001b[0m\n",
" \u001b[32m+\u001b[39m \u001b[1mhttpx\u001b[0m\u001b[2m==0.28.1\u001b[0m\n",
" \u001b[32m+\u001b[39m \u001b[1mhyperframe\u001b[0m\u001b[2m==6.1.0\u001b[0m\n",
" \u001b[32m+\u001b[39m \u001b[1midna\u001b[0m\u001b[2m==3.10\u001b[0m\n",
" \u001b[32m+\u001b[39m \u001b[1mimportlib-metadata\u001b[0m\u001b[2m==8.6.1\u001b[0m\n",
" \u001b[32m+\u001b[39m \u001b[1mimportlib-resources\u001b[0m\u001b[2m==6.5.2\u001b[0m\n",
" \u001b[32m+\u001b[39m \u001b[1mitsdangerous\u001b[0m\u001b[2m==2.2.0\u001b[0m\n",
" \u001b[32m+\u001b[39m \u001b[1mjinja2\u001b[0m\u001b[2m==3.1.6\u001b[0m\n",
" \u001b[32m+\u001b[39m \u001b[1mjiter\u001b[0m\u001b[2m==0.8.2\u001b[0m\n",
" \u001b[32m+\u001b[39m \u001b[1mjoblib\u001b[0m\u001b[2m==1.4.2\u001b[0m\n",
" \u001b[32m+\u001b[39m \u001b[1mkiwisolver\u001b[0m\u001b[2m==1.4.7\u001b[0m\n",
" \u001b[32m+\u001b[39m \u001b[1mmako\u001b[0m\u001b[2m==1.3.9\u001b[0m\n",
" \u001b[32m+\u001b[39m \u001b[1mmarkdown\u001b[0m\u001b[2m==3.7\u001b[0m\n",
" \u001b[32m+\u001b[39m \u001b[1mmarkupsafe\u001b[0m\u001b[2m==3.0.2\u001b[0m\n",
" \u001b[32m+\u001b[39m \u001b[1mmatplotlib\u001b[0m\u001b[2m==3.9.4\u001b[0m\n",
" \u001b[32m+\u001b[39m \u001b[1mmlflow\u001b[0m\u001b[2m==2.20.3\u001b[0m\n",
" \u001b[32m+\u001b[39m \u001b[1mmlflow-skinny\u001b[0m\u001b[2m==2.20.3\u001b[0m\n",
" \u001b[32m+\u001b[39m \u001b[1mnumpy\u001b[0m\u001b[2m==1.26.4\u001b[0m\n",
" \u001b[32m+\u001b[39m \u001b[1mopenai\u001b[0m\u001b[2m==1.63.0\u001b[0m\n",
" \u001b[32m+\u001b[39m \u001b[1mopentelemetry-api\u001b[0m\u001b[2m==1.16.0\u001b[0m\n",
" \u001b[32m+\u001b[39m \u001b[1mopentelemetry-sdk\u001b[0m\u001b[2m==1.16.0\u001b[0m\n",
" \u001b[32m+\u001b[39m \u001b[1mopentelemetry-semantic-conventions\u001b[0m\u001b[2m==0.37b0\u001b[0m\n",
" \u001b[32m+\u001b[39m \u001b[1mpackaging\u001b[0m\u001b[2m==24.2\u001b[0m\n",
" \u001b[32m+\u001b[39m \u001b[1mpandas\u001b[0m\u001b[2m==2.1.3\u001b[0m\n",
" \u001b[32m+\u001b[39m \u001b[1mpillow\u001b[0m\u001b[2m==11.1.0\u001b[0m\n",
" \u001b[32m+\u001b[39m \u001b[1mprotobuf\u001b[0m\u001b[2m==5.29.3\u001b[0m\n",
" \u001b[32m+\u001b[39m \u001b[1mpyarrow\u001b[0m\u001b[2m==19.0.1\u001b[0m\n",
" \u001b[32m+\u001b[39m \u001b[1mpyasn1\u001b[0m\u001b[2m==0.6.1\u001b[0m\n",
" \u001b[32m+\u001b[39m \u001b[1mpyasn1-modules\u001b[0m\u001b[2m==0.4.1\u001b[0m\n",
" \u001b[32m+\u001b[39m \u001b[1mpydantic\u001b[0m\u001b[2m==2.10.5\u001b[0m\n",
" \u001b[32m+\u001b[39m \u001b[1mpydantic-core\u001b[0m\u001b[2m==2.27.2\u001b[0m\n",
" \u001b[32m+\u001b[39m \u001b[1mpyparsing\u001b[0m\u001b[2m==3.2.1\u001b[0m\n",
" \u001b[32m+\u001b[39m \u001b[1mpython-dateutil\u001b[0m\u001b[2m==2.9.0.post0\u001b[0m\n",
" \u001b[32m+\u001b[39m \u001b[1mpytz\u001b[0m\u001b[2m==2025.1\u001b[0m\n",
" \u001b[32m+\u001b[39m \u001b[1mpyyaml\u001b[0m\u001b[2m==6.0.2\u001b[0m\n",
" \u001b[32m+\u001b[39m \u001b[1mrequests\u001b[0m\u001b[2m==2.32.3\u001b[0m\n",
" \u001b[32m+\u001b[39m \u001b[1mrsa\u001b[0m\u001b[2m==4.9\u001b[0m\n",
" \u001b[32m+\u001b[39m \u001b[1mscikit-learn\u001b[0m\u001b[2m==1.6.1\u001b[0m\n",
" \u001b[32m+\u001b[39m \u001b[1mscipy\u001b[0m\u001b[2m==1.13.1\u001b[0m\n",
" \u001b[32m+\u001b[39m \u001b[1msix\u001b[0m\u001b[2m==1.17.0\u001b[0m\n",
" \u001b[32m+\u001b[39m \u001b[1msmmap\u001b[0m\u001b[2m==5.0.2\u001b[0m\n",
" \u001b[32m+\u001b[39m \u001b[1msniffio\u001b[0m\u001b[2m==1.3.1\u001b[0m\n",
" \u001b[32m+\u001b[39m \u001b[1msqlalchemy\u001b[0m\u001b[2m==2.0.38\u001b[0m\n",
" \u001b[32m+\u001b[39m \u001b[1msqlparse\u001b[0m\u001b[2m==0.5.3\u001b[0m\n",
" \u001b[32m+\u001b[39m \u001b[1mthreadpoolctl\u001b[0m\u001b[2m==3.5.0\u001b[0m\n",
" \u001b[32m+\u001b[39m \u001b[1mtqdm\u001b[0m\u001b[2m==4.67.1\u001b[0m\n",
" \u001b[32m+\u001b[39m \u001b[1mtyping-extensions\u001b[0m\u001b[2m==4.12.2\u001b[0m\n",
" \u001b[32m+\u001b[39m \u001b[1mtzdata\u001b[0m\u001b[2m==2025.1\u001b[0m\n",
" \u001b[32m+\u001b[39m \u001b[1murllib3\u001b[0m\u001b[2m==2.3.0\u001b[0m\n",
" \u001b[32m+\u001b[39m \u001b[1mwerkzeug\u001b[0m\u001b[2m==3.1.3\u001b[0m\n",
" \u001b[32m+\u001b[39m \u001b[1mwrapt\u001b[0m\u001b[2m==1.17.2\u001b[0m\n",
" \u001b[32m+\u001b[39m \u001b[1mzipp\u001b[0m\u001b[2m==3.21.0\u001b[0m\n",
"2025/03/06 19:10:22 INFO mlflow.utils.environment: === Running command '['bash', '-c', 'source /var/folders/9g/psrbbvm92t712cy09d7_00d00000gp/T/tmpp18g94f9/envs/virtualenv_envs/mlflow-9d81fff15053e6e06e2edaefcc9e075d6c04a094/bin/activate && python -c \"\"']'\n",
"2025/03/06 19:10:22 INFO mlflow.utils.environment: === Running command '['bash', '-c', 'source /var/folders/9g/psrbbvm92t712cy09d7_00d00000gp/T/tmpp18g94f9/envs/virtualenv_envs/mlflow-9d81fff15053e6e06e2edaefcc9e075d6c04a094/bin/activate && python /Users/serena.ruan/Documents/repos/mlflow/mlflow/pyfunc/_mlflow_pyfunc_backend_predict.py --model-uri file:///Users/serena.ruan/Documents/test/mlruns/0/33b7da4d1693490b97934a5781964766/artifacts/model --content-type json --input-path /var/folders/9g/psrbbvm92t712cy09d7_00d00000gp/T/tmpz3dlhg3n/input.json']'\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"{\"predictions\": \"New York is a state located in the northeastern region of the United States. It is bordered by Vermont to the northeast, Massachusetts to the east, Connecticut to the southeast, and New Jersey and Pennsylvania to the south. The state also has access to the Atlantic Ocean to the southeast. The city of New York, often referred to simply as NYC, is the largest city in the state and is known for its significant cultural, financial, and historical influence.\"}"
]
}
],
"source": [
"# verify model before serving\n",
"mlflow.models.predict(\n",
" model_uri=model_info.model_uri,\n",
" input_data=[{\"role\": \"user\", \"content\": \"Where is New York?\"}],\n",
" env_manager=\"uv\",\n",
")"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "386566db-841a-41cf-a46f-d75bd060b1f8",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"runs:/33b7da4d1693490b97934a5781964766/model\n"
]
}
],
"source": [
"model_info.model_uri"
]
},
{
"cell_type": "markdown",
"id": "a4f72697-57c9-4a02-9c11-da35c277bab8",
"metadata": {},
"source": [
"## Serve the model locally"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "44e0e4e4-372b-4176-85e5-6f624b728c34",
"metadata": {},
"outputs": [],
"source": [
"# run below command to serve the model locally\n",
"# mlflow models serve -m runs:/33b7da4d1693490b97934a5781964766/model -p 6666"
]
},
{
"cell_type": "code",
"execution_count": 19,
"id": "85013e9c-fb42-4d8a-a8f1-5fb69f695b97",
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"b'{\"predictions\": \"British Shorthairs are generally considered to be intelligent cats, though their intelligence may manifest differently compared to some other breeds. They are known for their calm and laid-back demeanor, which can sometimes be mistaken for a lack of intelligence. In reality, they are capable of problem-solving and can be trained to perform basic commands or tricks, though they may not be as eager to please as some more active breeds.\\\\n\\\\nTheir intelligence is often reflected in their ability to adapt to their environment and their understanding of routines. British Shorthairs tend to be independent and may not seek out interaction as much as other, more playful breeds, but they can still form strong bonds with their owners. Essentially, while they might not be the most overtly intelligent cats, they possess a subtle understanding of their surroundings that reflects their adaptability and awareness.\"}'"
]
},
"execution_count": 19,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"import requests\n",
"\n",
"from mlflow.models.utils import convert_input_example_to_serving_input\n",
"\n",
"payload = convert_input_example_to_serving_input([\n",
" {\"role\": \"user\", \"content\": \"Is British shorthair smart?\"}\n",
"])\n",
"resp = requests.post(\n",
" \"http://127.0.0.1:6666/invocations\", data=payload, headers={\"Content-Type\": \"application/json\"}\n",
")\n",
"resp.content"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.10.15"
}
},
"nbformat": 4,
"nbformat_minor": 5
}
+5
View File
@@ -0,0 +1,5 @@
# MLflow Deployments
The examples provided within this directory show how to get started with MLflow Deployments using:
- Databricks (see the `databricks` subdirectory)
@@ -0,0 +1,112 @@
"""
Usage
-----
databricks secrets create-scope <scope>
databricks secrets put-secret <scope> openai-api-key --string-value $OPENAI_API_KEY
python examples/deployments/databricks.py --secret <scope>/openai-api-key
-----
"""
import argparse
import uuid
from mlflow.deployments import get_deploy_client
def parse_args():
parser = argparse.ArgumentParser()
parser.add_argument("--secret", type=str, help="Secret (e.g. secrets/scope/key)")
return parser.parse_args()
def main():
args = parse_args()
client = get_deploy_client("databricks")
name = f"test-endpoint-{uuid.uuid4()}"
client.create_endpoint(
name=name,
config={
"served_entities": [
{
"name": "test",
"external_model": {
"name": "gpt-4",
"provider": "openai",
"task": "llm/v1/chat",
"openai_config": {
"openai_api_key": "{{" + args.secret + "}}",
},
},
}
],
"tags": [
{
"key": "foo",
"value": "bar",
}
],
"rate_limits": [
{
"key": "user",
"renewal_period": "minute",
"calls": 5,
}
],
},
)
try:
# Update served_entities
print(
client.update_endpoint(
endpoint=name,
config={
"served_entities": [
{
"name": "test",
"external_model": {
"name": "gpt-4",
"provider": "openai",
"task": "llm/v1/chat",
"openai_config": {
"openai_api_key": "{{" + args.secret + "}}",
},
},
}
],
},
)
)
# Update rate_limits
print(
client.update_endpoint(
endpoint=name,
config={
"rate_limits": [
{
"key": "user",
"renewal_period": "minute",
"calls": 10,
}
],
},
)
)
print(client.list_endpoints()[:3])
print(client.get_endpoint(endpoint=name))
print(
client.predict(
endpoint=name,
inputs={
"messages": [
{"role": "user", "content": "Hello!"},
],
"max_tokens": 128,
},
),
)
finally:
client.delete_endpoint(endpoint=name)
if __name__ == "__main__":
main()
@@ -0,0 +1,113 @@
"""
Demo: MLflow Diffusers Adapter Flavor (LoRA)
This script demonstrates the full workflow of logging and loading a diffusion
model LoRA adapter using the native mlflow.diffusers flavor.
No GPU or real model weights required — uses a fake adapter for validation.
"""
import tempfile
from pathlib import Path
import numpy as np
import yaml
from safetensors.numpy import save_file
import mlflow
import mlflow.diffusers
def create_fake_lora_adapter(output_dir: Path) -> Path:
output_dir.mkdir(parents=True, exist_ok=True)
# Simulate LoRA weight matrices (small random tensors)
tensors = {
"unet.down_blocks.0.attentions.0.transformer_blocks.0.attn1.to_q.lora_down.weight": (
np.random.randn(4, 320).astype(np.float32)
),
"unet.down_blocks.0.attentions.0.transformer_blocks.0.attn1.to_q.lora_up.weight": (
np.random.randn(320, 4).astype(np.float32)
),
}
adapter_file = output_dir / "pytorch_lora_weights.safetensors"
save_file(tensors, str(adapter_file))
print(f"Created fake LoRA adapter at: {adapter_file}")
print(f" Adapter size: {adapter_file.stat().st_size} bytes")
return output_dir
def demo_log_and_load():
"""Demonstrate the full log -> load -> inspect cycle."""
with tempfile.TemporaryDirectory() as tmpdir:
# 1. Create fake adapter
adapter_dir = create_fake_lora_adapter(Path(tmpdir) / "my_lora")
# 2. Log the adapter with MLflow
print("\n--- Logging adapter with mlflow.diffusers.log_model() ---")
mlflow.set_experiment("diffusers-adapter-poc")
with mlflow.start_run(run_name="lora-adapter-demo") as run:
model_info = mlflow.diffusers.log_model(
adapter_path=str(adapter_dir),
base_model="black-forest-labs/FLUX.1-dev",
adapter_type="lora",
name="lora_model",
metadata={
"lora_rank": 4,
"training_steps": 1000,
"trigger_word": "sks style",
},
)
print(f" Run ID: {run.info.run_id}")
print(f" Model URI: {model_info.model_uri}")
# 3. Inspect the MLmodel file
print("\n--- MLmodel file contents ---")
model_uri = f"runs:/{run.info.run_id}/lora_model"
local_path = mlflow.artifacts.download_artifacts(model_uri)
mlmodel_path = Path(local_path) / "MLmodel"
with open(mlmodel_path) as f:
mlmodel = yaml.safe_load(f)
print(yaml.dump(mlmodel, default_flow_style=False, indent=2))
# 4. Load the model back
print("--- Loading model back with mlflow.diffusers.load_model() ---")
loaded = mlflow.diffusers.load_model(model_uri)
print(f" Type: {type(loaded).__name__}")
print(f" Base model: {loaded.base_model}")
print(f" Adapter type: {loaded.adapter_type}")
print(f" Adapter path: {loaded.adapter_path}")
print(f" Adapter files: {list(Path(loaded.adapter_path).iterdir())}")
# 5. Verify flavor config from MLmodel
print("\n--- Flavor config ---")
flavor_conf = mlmodel["flavors"]["diffusers"]
print(f" base_model: {flavor_conf['base_model']}")
print(f" adapter_type: {flavor_conf['adapter_type']}")
print(f" adapter_weights: {flavor_conf['adapter_weights']}")
# 6. Show that pyfunc interface is available
print("\n--- Pyfunc model interface ---")
print(" mlflow.pyfunc.load_model() would return a wrapper with predict()")
print(" predict() accepts: DataFrame/dict with 'prompt' column")
print(" predict() returns: list of PNG-encoded image bytes")
print(" (Skipping actual pyfunc load — requires base model download)")
print("\n--- Demo complete! ---")
print(
"The adapter is logged as a first-class MLflow model with full model registry support."
)
print(
"To generate images, call loaded.load_pipeline() on a machine "
"with the base model available."
)
if __name__ == "__main__":
demo_log_and_load()
View File
+6
View File
@@ -0,0 +1,6 @@
FROM python:3.8
RUN pip install mlflow azure-storage-blob numpy scipy pandas scikit-learn cloudpickle
COPY train.py .
COPY wine-quality.csv .
+11
View File
@@ -0,0 +1,11 @@
name: docker-example
docker_env:
image: mlflow-docker-example
entry_points:
main:
parameters:
alpha: float
l1_ratio: {type: float, default: 0.1}
command: "python train.py --alpha {alpha} --l1-ratio {l1_ratio}"
+71
View File
@@ -0,0 +1,71 @@
Dockerized Model Training with MLflow
-------------------------------------
This directory contains an MLflow project that trains a linear regression model on the UC Irvine
Wine Quality Dataset. The project uses a Docker image to capture the dependencies needed to run
training code. Running a project in a Docker environment (as opposed to Conda) allows for capturing
non-Python dependencies, e.g. Java libraries. In the future, we also hope to add tools to MLflow
for running Dockerized projects e.g. on a Kubernetes cluster for scale out.
Structure of this MLflow Project
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
This MLflow project contains a ``train.py`` file that trains a scikit-learn model and uses
MLflow Tracking APIs to log the model and its metadata (e.g., hyperparameters and metrics)
for later use and reference. ``train.py`` operates on the Wine Quality Dataset, which is included
in ``wine-quality.csv``.
Most importantly, the project also includes an ``MLproject`` file, which specifies the Docker
container environment in which to run the project using the ``docker_env`` field:
.. code-block:: yaml
docker_env:
image: mlflow-docker-example
Here, ``image`` can be any valid argument to ``docker run``, such as the tag, ID or URL of a Docker
image (see `Docker docs <https://docs.docker.com/engine/reference/run/#general-form>`_). The above
example references a locally-stored image (``mlflow-docker-example``) by tag.
Finally, the project includes a ``Dockerfile`` that is used to build the image referenced by the
``MLproject`` file. The ``Dockerfile`` specifies library dependencies required by the project, such
as ``mlflow`` and ``scikit-learn``.
Running this Example
^^^^^^^^^^^^^^^^^^^^
First, install MLflow (via ``pip install mlflow``) and install
`Docker <https://www.docker.com/get-started>`_.
Then, build the image for the project's Docker container environment. You must use the same image
name that is given by the ``docker_env.image`` field of the MLproject file. In this example, the
image name is ``mlflow-docker-example``. Issue the following command to build an image with this
name:
.. code-block:: bash
docker build -t mlflow-docker-example -f Dockerfile .
Note that the name if the image used in the ``docker build`` command, ``mlflow-docker-example``,
matches the name of the image referenced in the ``MLproject`` file.
Finally, run the example project using ``mlflow run examples/docker -P alpha=0.5``.
.. note::
If running this example on a Mac with Apple silicon, ensure that Docker Desktop is running and
that you are logged in to the Docker Desktop service.
If you are modifying the example ``DockerFile`` to specify older versions of ``scikit-learn``,
you should enable `Rosetta compatibility <https://docs.docker.com/desktop/settings/mac/#features-in-development>`_
in the Docker Desktop configuration settings to ensure that the appropriate ``cython`` compiler is used.
What happens when the project is run?
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
Running ``mlflow run examples/docker`` builds a new Docker image based on ``mlflow-docker-example``
that also contains our project code. The resulting image is tagged as
``mlflow-docker-example-<git-version>`` where ``<git-version>`` is the git commit ID. After the image is
built, MLflow executes the default (main) project entry point within the container using ``docker run``.
Environment variables, such as ``MLFLOW_TRACKING_URI``, are propagated inside the container during
project execution. When running against a local tracking URI, MLflow mounts the host system's
tracking directory (e.g., a local ``mlruns`` directory) inside the container so that metrics and
params logged during project execution are accessible afterwards.
+5
View File
@@ -0,0 +1,5 @@
{
"kube-context": "docker-for-desktop",
"kube-job-template-path": "examples/docker/kubernetes_job_template.yaml",
"repository-uri": "username/mlflow-kubernetes-example"
}
@@ -0,0 +1,20 @@
apiVersion: batch/v1
kind: Job
metadata:
name: "{replaced with MLflow Project name}"
namespace: mlflow
spec:
ttlSecondsAfterFinished: 100
backoffLimit: 0
template:
spec:
containers:
- name: "{replaced with MLflow Project name}"
image: "{replaced with URI of Docker image created during Project execution}"
command: ["{replaced with MLflow Project entry point command}"]
resources:
limits:
memory: 512Mi
requests:
memory: 256Mi
restartPolicy: Never
+70
View File
@@ -0,0 +1,70 @@
# The data set used in this example is from http://archive.ics.uci.edu/ml/datasets/Wine+Quality
# P. Cortez, A. Cerdeira, F. Almeida, T. Matos and J. Reis.
# Modeling wine preferences by data mining from physicochemical properties. In Decision Support Systems, Elsevier, 47(4):547-553, 2009.
import argparse
import os
import warnings
import numpy as np
import pandas as pd
from sklearn.linear_model import ElasticNet
from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score
from sklearn.model_selection import train_test_split
import mlflow
import mlflow.sklearn
def eval_metrics(actual, pred):
rmse = np.sqrt(mean_squared_error(actual, pred))
mae = mean_absolute_error(actual, pred)
r2 = r2_score(actual, pred)
return rmse, mae, r2
if __name__ == "__main__":
warnings.filterwarnings("ignore")
np.random.seed(40)
parser = argparse.ArgumentParser()
parser.add_argument("--alpha")
parser.add_argument("--l1-ratio")
args = parser.parse_args()
# Read the wine-quality csv file (make sure you're running this from the root of MLflow!)
wine_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "wine-quality.csv")
data = pd.read_csv(wine_path)
# Split the data into training and test sets. (0.75, 0.25) split.
train, test = train_test_split(data)
# The predicted column is "quality" which is a scalar from [3, 9]
train_x = train.drop(["quality"], axis=1)
test_x = test.drop(["quality"], axis=1)
train_y = train[["quality"]]
test_y = test[["quality"]]
alpha = float(args.alpha)
l1_ratio = float(args.l1_ratio)
with mlflow.start_run():
lr = ElasticNet(alpha=alpha, l1_ratio=l1_ratio, random_state=42)
lr.fit(train_x, train_y)
predicted_qualities = lr.predict(test_x)
(rmse, mae, r2) = eval_metrics(test_y, predicted_qualities)
print(f"Elasticnet model (alpha={alpha:f}, l1_ratio={l1_ratio:f}):")
print(f" RMSE: {rmse}")
print(f" MAE: {mae}")
print(f" R2: {r2}")
mlflow.log_param("alpha", alpha)
mlflow.log_param("l1_ratio", l1_ratio)
mlflow.log_metric("rmse", rmse)
mlflow.log_metric("r2", r2)
mlflow.log_metric("mae", mae)
mlflow.sklearn.log_model(lr, name="model")
File diff suppressed because it is too large Load Diff
+40
View File
@@ -0,0 +1,40 @@
### MLflow evaluation Examples
The examples in this directory demonstrate how to use the `mlflow.evaluate()` API. Specifically,
they show how to evaluate a PyFunc model on a specified dataset using the builtin default evaluator
and specified extra metrics, where the resulting metrics & artifacts are logged to MLflow Tracking.
They also show how to specify validation thresholds for the resulting metrics to validate the quality
of your model. See full list of examples below:
- Example `evaluate_on_binary_classifier.py` evaluates an xgboost `XGBClassifier` model on dataset loaded by
`shap.datasets.adult`.
- Example `evaluate_on_multiclass_classifier.py` evaluates a scikit-learn `LogisticRegression` model on dataset
generated by `sklearn.datasets.make_classification`.
- Example `evaluate_on_regressor.py` evaluate as scikit-learn `LinearRegression` model on dataset loaded by
`sklearn.datasets.load_diabetes`
- Example `evaluate_with_custom_metrics.py` evaluates a scikit-learn `LinearRegression`
model with a custom metric function on dataset loaded by `sklearn.datasets.load_diabetes`
- Example `evaluate_with_custom_metrics_comprehensive.py` evaluates a scikit-learn `LinearRegression` model
with a comprehensive list of custom metric functions on dataset loaded by `sklearn.datasets.load_diabetes`
- Example `evaluate_with_model_validation.py` trains both a candidate xgboost `XGBClassifier` model
and a baseline `DummyClassifier` model on dataset loaded by `shap.datasets.adult`. Then, it validates
the candidate model against specified thresholds on both builtin and extra metrics and the dummy model.
#### Prerequisites
```
pip install scikit-learn xgboost shap>=0.40 matplotlib
```
#### How to run the examples
Run in this directory with Python.
```sh
python evaluate_on_binary_classifier.py
python evaluate_on_multiclass_classifier.py
python evaluate_on_regressor.py
python evaluate_with_custom_metrics.py
python evaluate_with_custom_metrics_comprehensive.py
python evaluate_with_model_validation.py
```
@@ -0,0 +1,41 @@
import shap
import xgboost
from sklearn.model_selection import train_test_split
import mlflow
from mlflow.models import infer_signature
# Load the UCI Adult Dataset
X, y = shap.datasets.adult()
# Split the data into training and test sets
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.33, random_state=42)
# Fit an XGBoost binary classifier on the training data split
model = xgboost.XGBClassifier().fit(X_train, y_train)
# Infer model signature
predictions = model.predict(X_train)
signature = infer_signature(X_train, predictions)
# Build the Evaluation Dataset from the test set
eval_data = X_test
eval_data["label"] = y_test
with mlflow.start_run() as run:
# Log the XGBoost binary classifier model to MLflow
model_info = mlflow.sklearn.log_model(
model, name="model", signature=signature, serialization_format="cloudpickle"
)
# Evaluate the logged model
result = mlflow.evaluate(
model_info.model_uri,
eval_data,
targets="label",
model_type="classifier",
evaluators=["default"],
)
print(f"metrics:\n{result.metrics}")
print(f"artifacts:\n{result.artifacts}")
@@ -0,0 +1,25 @@
from sklearn.datasets import make_classification
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import train_test_split
import mlflow
X, y = make_classification(n_samples=10000, n_classes=10, n_informative=5, random_state=1)
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.33, random_state=42)
with mlflow.start_run() as run:
model = LogisticRegression(solver="liblinear").fit(X_train, y_train)
model_info = mlflow.sklearn.log_model(model, name="model")
result = mlflow.evaluate(
model_info.model_uri,
X_test,
targets=y_test,
model_type="classifier",
evaluators="default",
evaluator_config={"log_model_explainability": True, "explainability_nsamples": 1000},
)
print(f"run_id={run.info.run_id}")
print(f"metrics:\n{result.metrics}")
print(f"artifacts:\n{result.artifacts}")
@@ -0,0 +1,28 @@
from sklearn.datasets import load_diabetes
from sklearn.linear_model import LinearRegression
from sklearn.model_selection import train_test_split
import mlflow
diabetes_dataset = load_diabetes()
X_train, X_test, y_train, y_test = train_test_split(
diabetes_dataset.data, diabetes_dataset.target, test_size=0.33, random_state=42
)
with mlflow.start_run() as run:
model = LinearRegression().fit(X_train, y_train)
model_info = mlflow.sklearn.log_model(model, name="model")
result = mlflow.evaluate(
model_info.model_uri,
X_test,
targets=y_test,
model_type="regressor",
evaluators="default",
feature_names=diabetes_dataset.feature_names,
evaluator_config={"explainability_nsamples": 1000},
)
print(f"metrics:\n{result.metrics}")
print(f"artifacts:\n{result.artifacts}")
@@ -0,0 +1,67 @@
import os
import openai
import pandas as pd
import mlflow
from mlflow.metrics import make_metric
from mlflow.metrics.base import MetricValue, standard_aggregations
assert "OPENAI_API_KEY" in os.environ, "Please set the OPENAI_API_KEY environment variable."
# Helper function to check if a string is valid python code
def is_valid_python_code(code: str) -> bool:
try:
compile(code, "<string>", "exec")
return True
except SyntaxError:
return False
# Create an evaluation function that iterates through the predictions
def eval_fn(predictions):
scores = [int(is_valid_python_code(prediction)) for prediction in predictions]
return MetricValue(
scores=scores,
aggregate_results=standard_aggregations(scores),
)
# Create an EvaluationMetric object for the python code metric
valid_code_metric = make_metric(
eval_fn=eval_fn, greater_is_better=False, name="valid_python_code", version="v1"
)
eval_df = pd.DataFrame({
"input": [
"SELECT * FROM ",
"import pandas",
"def hello_world",
],
})
with mlflow.start_run() as run:
system_prompt = (
"Generate code that is less than 50 characters. Return only python code and nothing else."
)
logged_model = mlflow.openai.log_model(
model="gpt-4o-mini",
task=openai.chat.completions,
name="model",
messages=[
{"role": "system", "content": system_prompt},
{"role": "user", "content": "{question}"},
],
)
results = mlflow.evaluate(
logged_model.model_uri,
eval_df,
model_type="text",
extra_metrics=[valid_code_metric],
)
print(results)
eval_table = results.tables["eval_results_table"]
print(eval_table)
@@ -0,0 +1,85 @@
import os
import matplotlib.pyplot as plt
import numpy as np
from sklearn.datasets import load_diabetes
from sklearn.linear_model import LinearRegression
from sklearn.model_selection import train_test_split
import mlflow
from mlflow.models import infer_signature, make_metric
# loading the diabetes dataset
diabetes_dataset = load_diabetes(as_frame=True)
# split the dataset into train and test partitions
X_train, X_test, y_train, y_test = train_test_split(
diabetes_dataset.data, diabetes_dataset.target, test_size=0.2, random_state=123
)
# train the model
lin_reg = LinearRegression().fit(X_train, y_train)
# Infer model signature
predictions = lin_reg.predict(X_train)
signature = infer_signature(X_train, predictions)
# creating the evaluation dataframe
eval_data = X_test.copy()
eval_data["target"] = y_test
def squared_diff_plus_one(eval_df, _builtin_metrics):
"""
This example custom metric function creates a metric based on the ``prediction`` and
``target`` columns in ``eval_df`.
"""
return np.sum(np.abs(eval_df["prediction"] - eval_df["target"] + 1) ** 2)
def sum_on_target_divided_by_two(_eval_df, builtin_metrics):
"""
This example custom metric function creates a metric derived from existing metrics in
``builtin_metrics``.
"""
return builtin_metrics["sum_on_target"] / 2
def prediction_target_scatter(eval_df, _builtin_metrics, artifacts_dir):
"""
This example custom artifact generates and saves a scatter plot to ``artifacts_dir`` that
visualizes the relationship between the predictions and targets for the given model to a
file as an image artifact.
"""
plt.scatter(eval_df["prediction"], eval_df["target"])
plt.xlabel("Targets")
plt.ylabel("Predictions")
plt.title("Targets vs. Predictions")
plot_path = os.path.join(artifacts_dir, "example_scatter_plot.png")
plt.savefig(plot_path)
return {"example_scatter_plot_artifact": plot_path}
with mlflow.start_run() as run:
model_info = mlflow.sklearn.log_model(lin_reg, name="model", signature=signature)
result = mlflow.evaluate(
model=model_info.model_uri,
data=eval_data,
targets="target",
model_type="regressor",
evaluators=["default"],
extra_metrics=[
make_metric(
eval_fn=squared_diff_plus_one,
greater_is_better=False,
),
make_metric(
eval_fn=sum_on_target_divided_by_two,
greater_is_better=True,
),
],
custom_artifacts=[prediction_target_scatter],
)
print(f"metrics:\n{result.metrics}")
print(f"artifacts:\n{result.artifacts}")
@@ -0,0 +1,84 @@
import numpy as np
import pandas as pd
from matplotlib.figure import Figure
from sklearn.datasets import load_diabetes
from sklearn.linear_model import LinearRegression
from sklearn.model_selection import train_test_split
import mlflow
from mlflow.models import infer_signature, make_metric
# loading the diabetes dataset
diabetes_dataset = load_diabetes(as_frame=True)
# split the dataset into train and test partitions
X_train, X_test, y_train, y_test = train_test_split(
diabetes_dataset.data, diabetes_dataset.target, test_size=0.2, random_state=123
)
# train the model
lin_reg = LinearRegression().fit(X_train, y_train)
# Infer model signature
predictions = lin_reg.predict(X_train)
signature = infer_signature(X_train, predictions)
# creating the evaluation dataframe
eval_data = X_test.copy()
eval_data["target"] = y_test
def custom_metric(eval_df, _builtin_metrics):
return np.sum(np.abs(eval_df["prediction"] - eval_df["target"] + 1) ** 2)
class ExampleClass:
def __init__(self, x):
self.x = x
def custom_artifact(eval_df, builtin_metrics, _artifacts_dir):
example_np_arr = np.array([1, 2, 3])
example_df = pd.DataFrame({"test": [2.2, 3.1], "test2": [3, 2]})
example_dict = {"hello": "there", "test_list": [0.1, 0.3, 4]}
example_dict.update(builtin_metrics)
example_dict_2 = '{"a": 3, "b": [1, 2, 3]}'
example_image = Figure()
ax = example_image.subplots()
ax.scatter(eval_df["prediction"], eval_df["target"])
ax.set_xlabel("Targets")
ax.set_ylabel("Predictions")
ax.set_title("Targets vs. Predictions")
example_custom_class = ExampleClass(10)
return {
"example_np_arr_from_obj_saved_as_npy": example_np_arr,
"example_df_from_obj_saved_as_csv": example_df,
"example_dict_from_obj_saved_as_json": example_dict,
"example_image_from_obj_saved_as_png": example_image,
"example_dict_from_json_str_saved_as_json": example_dict_2,
"example_class_from_obj_saved_as_pickle": example_custom_class,
}
with mlflow.start_run() as run:
model_info = mlflow.sklearn.log_model(lin_reg, name="model", signature=signature)
result = mlflow.evaluate(
model=model_info.model_uri,
data=eval_data,
targets="target",
model_type="regressor",
evaluators=["default"],
extra_metrics=[
make_metric(
eval_fn=custom_metric,
greater_is_better=False,
)
],
custom_artifacts=[
custom_artifact,
],
)
print(f"metrics:\n{result.metrics}")
print(f"artifacts:\n{result.artifacts}")
@@ -0,0 +1,37 @@
import shap
import xgboost
from sklearn.model_selection import train_test_split
import mlflow
# Load the UCI Adult Dataset
X, y = shap.datasets.adult()
# Split the data into training and test sets
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.33, random_state=42)
# Fit an XGBoost binary classifier on the training data split
model = xgboost.XGBClassifier().fit(X_train, y_train)
# Build the Evaluation Dataset from the test set
eval_data = X_test
eval_data["label"] = y_test
# Define a function that calls the model's predict method
def fn(X):
return model.predict(X)
with mlflow.start_run() as run:
# Evaluate the function without logging the model
result = mlflow.evaluate(
fn,
eval_data,
targets="label",
model_type="classifier",
evaluators=["default"],
)
print(f"metrics:\n{result.metrics}")
print(f"artifacts:\n{result.artifacts}")
@@ -0,0 +1,67 @@
import os
import openai
import pandas as pd
import mlflow
from mlflow.metrics.genai import EvaluationExample, answer_similarity
assert "OPENAI_API_KEY" in os.environ, "Please set the OPENAI_API_KEY environment variable."
# testing with OpenAI gpt-4o-mini
example = EvaluationExample(
input="What is MLflow?",
output="MLflow is an open-source platform for managing machine "
"learning workflows, including experiment tracking, model packaging, "
"versioning, and deployment, simplifying the ML lifecycle.",
score=4,
justification="The definition effectively explains what MLflow is "
"its purpose, and its developer. It could be more concise for a 5-score.",
grading_context={
"ground_truth": "MLflow is an open-source platform for managing "
"the end-to-end machine learning (ML) lifecycle. It was developed by Databricks, "
"a company that specializes in big data and machine learning solutions. MLflow is "
"designed to address the challenges that data scientists and machine learning "
"engineers face when developing, training, and deploying machine learning models."
},
)
answer_similarity_metric = answer_similarity(examples=[example])
eval_df = pd.DataFrame({
"inputs": [
"What is MLflow?",
"What is Spark?",
"What is Python?",
],
"ground_truth": [
"MLflow is an open-source platform for managing the end-to-end machine learning (ML) lifecycle. It was developed by Databricks, a company that specializes in big data and machine learning solutions. MLflow is designed to address the challenges that data scientists and machine learning engineers face when developing, training, and deploying machine learning models.",
"Apache Spark is an open-source, distributed computing system designed for big data processing and analytics. It was developed in response to limitations of the Hadoop MapReduce computing model, offering improvements in speed and ease of use. Spark provides libraries for various tasks such as data ingestion, processing, and analysis through its components like Spark SQL for structured data, Spark Streaming for real-time data processing, and MLlib for machine learning tasks",
"Python is a high-level programming language that was created by Guido van Rossum and released in 1991. It emphasizes code readability and allows developers to express concepts in fewer lines of code than languages like C++ or Java. Python is used in various domains, including web development, scientific computing, data analysis, and machine learning.",
],
})
with mlflow.start_run() as run:
system_prompt = "Answer the following question in two sentences"
logged_model = mlflow.openai.log_model(
model="gpt-4o-mini",
task=openai.chat.completions,
name="model",
messages=[
{"role": "system", "content": system_prompt},
{"role": "user", "content": "{question}"},
],
)
results = mlflow.evaluate(
logged_model.model_uri,
eval_df,
targets="ground_truth",
model_type="question-answering",
extra_metrics=[answer_similarity_metric],
)
print(results)
eval_table = results.tables["eval_results_table"]
print(eval_table)
@@ -0,0 +1,110 @@
import shap
import xgboost
from sklearn.dummy import DummyClassifier
from sklearn.model_selection import train_test_split
import mlflow
from mlflow.models import MetricThreshold, infer_signature, make_metric
# load UCI Adult Data Set; segment it into training and test sets
X, y = shap.datasets.adult()
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.33, random_state=42)
# train a candidate XGBoost model
candidate_model = xgboost.XGBClassifier().fit(X_train, y_train)
candidate_signature = infer_signature(X_train, candidate_model.predict(X_train))
# train a baseline dummy model
baseline_model = DummyClassifier(strategy="uniform").fit(X_train, y_train)
baseline_signature = infer_signature(X_train, baseline_model.predict(X_train))
# construct an evaluation dataset from the test set
eval_data = X_test
eval_data["label"] = y_test
# Define a custom metric to evaluate against
def double_positive(_eval_df, builtin_metrics):
return builtin_metrics["true_positives"] * 2
# Define criteria for model to be validated against
thresholds = {
# Specify metric value threshold
"precision_score": MetricThreshold(
threshold=0.7, greater_is_better=True
), # precision should be >=0.7
# Specify model comparison thresholds
"recall_score": MetricThreshold(
min_absolute_change=0.1, # recall should be at least 0.1 greater than baseline model recall
min_relative_change=0.1, # recall should be at least 10 percent greater than baseline model recall
greater_is_better=True,
),
# Specify both metric value and model comparison thresholds
"accuracy_score": MetricThreshold(
threshold=0.8, # accuracy should be >=0.8
min_absolute_change=0.05, # accuracy should be at least 0.05 greater than baseline model accuracy
min_relative_change=0.05, # accuracy should be at least 5 percent greater than baseline model accuracy
greater_is_better=True,
),
# Specify threshold for custom metric
"double_positive": MetricThreshold(
threshold=1e5,
greater_is_better=False, # double_positive should be <=1e5
),
}
double_positive_metric = make_metric(
eval_fn=double_positive,
greater_is_better=False,
)
with mlflow.start_run() as run:
# Note: in most model validation use-cases the baseline model should instead b
# a previously trained model (such as the current production model)
baseline_model_uri = mlflow.sklearn.log_model(
baseline_model,
name="baseline_model",
signature=baseline_signature,
serialization_format="cloudpickle",
).model_uri
# Evaluate the baseline model
baseline_result = mlflow.evaluate(
baseline_model_uri,
eval_data,
targets="label",
model_type="classifier",
extra_metrics=[double_positive_metric],
# set to env_manager to "virtualenv" or "conda" to score the candidate and baseline models
# in isolated Python environments where their dependencies are restored.
env_manager="local",
)
# Evaluate the candidate model
candidate_model_uri = mlflow.sklearn.log_model(
candidate_model,
name="candidate_model",
signature=candidate_signature,
serialization_format="cloudpickle",
).model_uri
candidate_result = mlflow.evaluate(
candidate_model_uri,
eval_data,
targets="label",
model_type="classifier",
extra_metrics=[double_positive_metric],
env_manager="local",
)
# Validate the candidate result against the baseline
mlflow.validate_evaluation_results(
candidate_result=candidate_result,
baseline_result=baseline_result,
validation_thresholds=thresholds,
)
# If you would like to catch model validation failures, you can add try except clauses around
# the mlflow.evaluate() call and catch the ModelValidationFailedException, imported at the top
# of this file.
@@ -0,0 +1,41 @@
import openai
import pandas as pd
import mlflow
eval_df = pd.DataFrame({
"inputs": [
"What is MLflow?",
"What is Spark?",
"What is Python?",
],
"ground_truth": [
"MLflow is an open-source platform for managing the end-to-end machine learning (ML) lifecycle. It was developed by Databricks, a company that specializes in big data and machine learning solutions. MLflow is designed to address the challenges that data scientists and machine learning engineers face when developing, training, and deploying machine learning models.",
"Apache Spark is an open-source, distributed computing system designed for big data processing and analytics. It was developed in response to limitations of the Hadoop MapReduce computing model, offering improvements in speed and ease of use. Spark provides libraries for various tasks such as data ingestion, processing, and analysis through its components like Spark SQL for structured data, Spark Streaming for real-time data processing, and MLlib for machine learning tasks",
"Python is a high-level programming language that was created by Guido van Rossum and released in 1991. It emphasizes code readability and allows developers to express concepts in fewer lines of code than languages like C++ or Java. Python is used in various domains, including web development, scientific computing, data analysis, and machine learning.",
],
})
with mlflow.start_run() as run:
system_prompt = "Answer the following question in two sentences"
logged_model = mlflow.openai.log_model(
model="gpt-4o-mini",
task=openai.chat.completions,
name="model",
messages=[
{"role": "system", "content": system_prompt},
{"role": "user", "content": "{question}"},
],
)
results = mlflow.evaluate(
logged_model.model_uri,
eval_df,
targets="ground_truth",
model_type="question-answering",
evaluators="default",
)
print(results.metrics)
eval_table = results.tables["eval_results_table"]
print(eval_table)
@@ -0,0 +1,33 @@
import shap
import xgboost
from sklearn.model_selection import train_test_split
import mlflow
# Load the UCI Adult Dataset
X, y = shap.datasets.adult()
# Split the data into training and test sets
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.33, random_state=42)
# Fit an XGBoost binary classifier on the training data split
model = xgboost.XGBClassifier().fit(X_train, y_train)
# Build the Evaluation Dataset from the test set
y_test_pred = model.predict(X=X_test)
eval_data = X_test
eval_data["label"] = y_test
eval_data["predictions"] = y_test_pred
with mlflow.start_run() as run:
# Evaluate the static dataset without providing a model
result = mlflow.evaluate(
data=eval_data,
targets="label",
predictions="predictions",
model_type="classifier",
)
print(f"metrics:\n{result.metrics}")
print(f"artifacts:\n{result.artifacts}")
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff
+590
View File
@@ -0,0 +1,590 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {
"application/vnd.databricks.v1+cell": {
"cellMetadata": {},
"inputWidgets": {},
"nuid": "42084110-295b-493a-9b3e-5d8d29ff78b3",
"showTitle": false,
"title": ""
}
},
"source": [
"# LLM RAG Evaluation with MLflow Example Notebook\n",
"\n",
"In this notebook, we will demonstrate how to evaluate various a RAG system with MLflow."
]
},
{
"cell_type": "markdown",
"metadata": {
"application/vnd.databricks.v1+cell": {
"cellMetadata": {},
"inputWidgets": {},
"nuid": "bdff35e3-0e09-48b8-87ce-78759de88998",
"showTitle": false,
"title": ""
}
},
"source": [
"We need to set our OpenAI API key, since we will be using GPT-4 for our LLM-judged metrics.\n",
"\n",
"In order to set your private key safely, please be sure to either export your key through a command-line terminal for your current instance, or, for a permanent addition to all user-based sessions, configure your favored environment management configuration file (i.e., .bashrc, .zshrc) to have the following entry:\n",
"\n",
"`OPENAI_API_KEY=<your openai API key>`"
]
},
{
"cell_type": "code",
"execution_count": 4,
"metadata": {
"application/vnd.databricks.v1+cell": {
"cellMetadata": {
"byteLimit": 2048000,
"rowLimit": 10000
},
"inputWidgets": {},
"nuid": "fb946228-62fb-4d68-9732-75935c9cb401",
"showTitle": false,
"title": ""
}
},
"outputs": [],
"source": [
"import pandas as pd\n",
"\n",
"import mlflow"
]
},
{
"cell_type": "markdown",
"metadata": {
"application/vnd.databricks.v1+cell": {
"cellMetadata": {},
"inputWidgets": {},
"nuid": "273d1345-95d7-435a-a7b6-a5f3dbb3f073",
"showTitle": false,
"title": ""
}
},
"source": [
"## Create a RAG system\n",
"\n",
"Use Langchain and Chroma to create a RAG system that answers questions based on the MLflow documentation."
]
},
{
"cell_type": "code",
"execution_count": 5,
"metadata": {
"application/vnd.databricks.v1+cell": {
"cellMetadata": {
"byteLimit": 2048000,
"rowLimit": 10000
},
"inputWidgets": {},
"nuid": "2c28d0ad-f469-46ab-a2b4-c5e8db50a729",
"showTitle": false,
"title": ""
}
},
"outputs": [],
"source": [
"from langchain.chains import RetrievalQA\n",
"from langchain.document_loaders import WebBaseLoader\n",
"from langchain.embeddings.openai import OpenAIEmbeddings\n",
"from langchain.llms import OpenAI\n",
"from langchain.text_splitter import CharacterTextSplitter\n",
"from langchain.vectorstores import Chroma"
]
},
{
"cell_type": "code",
"execution_count": 9,
"metadata": {
"application/vnd.databricks.v1+cell": {
"cellMetadata": {
"byteLimit": 2048000,
"rowLimit": 10000
},
"inputWidgets": {},
"nuid": "83a7e77e-6717-472a-86dc-02e2c356ddef",
"showTitle": false,
"title": ""
}
},
"outputs": [],
"source": [
"loader = WebBaseLoader(\"https://mlflow.org/docs/latest/index.html\")\n",
"\n",
"documents = loader.load()\n",
"text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)\n",
"texts = text_splitter.split_documents(documents)\n",
"\n",
"embeddings = OpenAIEmbeddings()\n",
"docsearch = Chroma.from_documents(texts, embeddings)\n",
"\n",
"qa = RetrievalQA.from_chain_type(\n",
" llm=OpenAI(temperature=0),\n",
" chain_type=\"stuff\",\n",
" retriever=docsearch.as_retriever(),\n",
" return_source_documents=True,\n",
")"
]
},
{
"cell_type": "markdown",
"metadata": {
"application/vnd.databricks.v1+cell": {
"cellMetadata": {},
"inputWidgets": {},
"nuid": "fd70bcf6-7c44-44d3-9435-567b82611e1c",
"showTitle": false,
"title": ""
}
},
"source": [
"## Evaluate the RAG system using `mlflow.evaluate()`"
]
},
{
"cell_type": "markdown",
"metadata": {
"application/vnd.databricks.v1+cell": {
"cellMetadata": {},
"inputWidgets": {},
"nuid": "de1bc359-2e40-459c-bea4-bed35a117988",
"showTitle": false,
"title": ""
}
},
"source": [
"Create a simple function that runs each input through the RAG chain"
]
},
{
"cell_type": "code",
"execution_count": 10,
"metadata": {
"application/vnd.databricks.v1+cell": {
"cellMetadata": {
"byteLimit": 2048000,
"rowLimit": 10000
},
"inputWidgets": {},
"nuid": "667ec809-2bb5-4170-9937-6804386b41ec",
"showTitle": false,
"title": ""
}
},
"outputs": [],
"source": [
"def model(input_df):\n",
" answer = []\n",
" for index, row in input_df.iterrows():\n",
" answer.append(qa(row[\"questions\"]))\n",
"\n",
" return answer"
]
},
{
"cell_type": "markdown",
"metadata": {
"application/vnd.databricks.v1+cell": {
"cellMetadata": {},
"inputWidgets": {},
"nuid": "d1064306-b7f3-4b3e-825c-4353d808f21d",
"showTitle": false,
"title": ""
}
},
"source": [
"Create an eval dataset"
]
},
{
"cell_type": "code",
"execution_count": 16,
"metadata": {
"application/vnd.databricks.v1+cell": {
"cellMetadata": {
"byteLimit": 2048000,
"rowLimit": 10000
},
"inputWidgets": {},
"nuid": "a5481491-e4a9-42ea-8a3f-f527faffd04d",
"showTitle": false,
"title": ""
}
},
"outputs": [],
"source": [
"eval_df = pd.DataFrame({\n",
" \"questions\": [\n",
" \"What is MLflow?\",\n",
" \"How to run mlflow.evaluate()?\",\n",
" \"How to log_table()?\",\n",
" \"How to load_table()?\",\n",
" ],\n",
"})"
]
},
{
"cell_type": "markdown",
"metadata": {
"application/vnd.databricks.v1+cell": {
"cellMetadata": {},
"inputWidgets": {},
"nuid": "9c3c8023-8feb-427a-b36d-34cd1853a5dc",
"showTitle": false,
"title": ""
}
},
"source": [
"Create a faithfulness metric"
]
},
{
"cell_type": "code",
"execution_count": 12,
"metadata": {
"application/vnd.databricks.v1+cell": {
"cellMetadata": {
"byteLimit": 2048000,
"rowLimit": 10000
},
"inputWidgets": {},
"nuid": "3882b940-9c25-41ce-a301-72d8c0c90aaa",
"showTitle": false,
"title": ""
}
},
"outputs": [],
"source": [
"from mlflow.metrics.genai.metric_definitions import faithfulness\n",
"\n",
"faithfulness_metric = faithfulness(model=\"openai:/gpt-4\")"
]
},
{
"cell_type": "code",
"execution_count": 17,
"metadata": {
"application/vnd.databricks.v1+cell": {
"cellMetadata": {
"byteLimit": 2048000,
"rowLimit": 10000
},
"inputWidgets": {},
"nuid": "ea40ce52-6ac7-4c20-9669-d24f80a6cebe",
"showTitle": false,
"title": ""
}
},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"2023/10/23 13:13:16 INFO mlflow.models.evaluation.base: Evaluating the model with the default evaluator.\n"
]
},
{
"name": "stderr",
"output_type": "stream",
"text": [
"Number of requested results 4 is greater than number of elements in index 3, updating n_results = 3\n",
"Number of requested results 4 is greater than number of elements in index 3, updating n_results = 3\n",
"Number of requested results 4 is greater than number of elements in index 3, updating n_results = 3\n",
"Number of requested results 4 is greater than number of elements in index 3, updating n_results = 3\n",
"Using pad_token, but it is not set yet.\n"
]
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "23e9a5f58f1b4930ac47c88259156e1d",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
" 0%| | 0/1 [00:00<?, ?it/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"name": "stderr",
"output_type": "stream",
"text": [
"2023/10/23 13:13:41 INFO mlflow.models.evaluation.default_evaluator: Evaluating builtin metrics: token_count\n",
"2023/10/23 13:13:41 INFO mlflow.models.evaluation.default_evaluator: Evaluating builtin metrics: toxicity\n",
"2023/10/23 13:13:41 INFO mlflow.models.evaluation.default_evaluator: Evaluating builtin metrics: perplexity\n",
"Using pad_token, but it is not set yet.\n"
]
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "2c6fd2067bad4404ad5550d56e23407e",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
" 0%| | 0/1 [00:00<?, ?it/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"name": "stderr",
"output_type": "stream",
"text": [
"2023/10/23 13:13:44 INFO mlflow.models.evaluation.default_evaluator: Evaluating builtin metrics: flesch_kincaid_grade_level\n",
"2023/10/23 13:13:44 INFO mlflow.models.evaluation.default_evaluator: Evaluating builtin metrics: ari_grade_level\n",
"2023/10/23 13:13:44 INFO mlflow.models.evaluation.default_evaluator: Evaluating builtin metrics: exact_match\n",
"2023/10/23 13:13:44 INFO mlflow.models.evaluation.default_evaluator: Evaluating metrics: faithfulness\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"{'toxicity/v1/mean': 0.0002736186215770431, 'toxicity/v1/variance': 2.856656765360073e-08, 'toxicity/v1/p90': 0.0004570253004203551, 'toxicity/v1/ratio': 0.0, 'perplexity/v1/mean': 70.08646988868713, 'perplexity/v1/variance': 5233.465638493719, 'perplexity/v1/p90': 149.10144042968753, 'flesch_kincaid_grade_level/v1/mean': 7.625, 'flesch_kincaid_grade_level/v1/variance': 23.836875, 'flesch_kincaid_grade_level/v1/p90': 13.150000000000002, 'ari_grade_level/v1/mean': 9.450000000000001, 'ari_grade_level/v1/variance': 32.262499999999996, 'ari_grade_level/v1/p90': 15.870000000000001, 'faithfulness/v1/mean': 4.0, 'faithfulness/v1/variance': 3.0, 'faithfulness/v1/p90': 5.0}\n"
]
}
],
"source": [
"results = mlflow.evaluate(\n",
" model,\n",
" eval_df,\n",
" model_type=\"question-answering\",\n",
" evaluators=\"default\",\n",
" predictions=\"result\",\n",
" extra_metrics=[faithfulness_metric, mlflow.metrics.latency()],\n",
" evaluator_config={\n",
" \"col_mapping\": {\n",
" \"inputs\": \"questions\",\n",
" \"context\": \"source_documents\",\n",
" }\n",
" },\n",
")\n",
"print(results.metrics)"
]
},
{
"cell_type": "code",
"execution_count": 18,
"metadata": {
"application/vnd.databricks.v1+cell": {
"cellMetadata": {},
"inputWidgets": {},
"nuid": "989a0861-5153-44e6-a19d-efcae7fe6cb5",
"showTitle": false,
"title": ""
}
},
"outputs": [
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "4a4883be06c94983a171da51d14b40a3",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"Downloading artifacts: 0%| | 0/1 [00:00<?, ?it/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"text/html": [
"<div>\n",
"<style scoped>\n",
" .dataframe tbody tr th:only-of-type {\n",
" vertical-align: middle;\n",
" }\n",
"\n",
" .dataframe tbody tr th {\n",
" vertical-align: top;\n",
" }\n",
"\n",
" .dataframe thead th {\n",
" text-align: right;\n",
" }\n",
"</style>\n",
"<table border=\"1\" class=\"dataframe\">\n",
" <thead>\n",
" <tr style=\"text-align: right;\">\n",
" <th></th>\n",
" <th>questions</th>\n",
" <th>outputs</th>\n",
" <th>query</th>\n",
" <th>source_documents</th>\n",
" <th>latency</th>\n",
" <th>token_count</th>\n",
" <th>toxicity/v1/score</th>\n",
" <th>perplexity/v1/score</th>\n",
" <th>flesch_kincaid_grade_level/v1/score</th>\n",
" <th>ari_grade_level/v1/score</th>\n",
" <th>faithfulness/v1/score</th>\n",
" <th>faithfulness/v1/justification</th>\n",
" </tr>\n",
" </thead>\n",
" <tbody>\n",
" <tr>\n",
" <th>0</th>\n",
" <td>What is MLflow?</td>\n",
" <td>MLflow is an open source platform for managin...</td>\n",
" <td>What is MLflow?</td>\n",
" <td>[{'lc_attributes': {}, 'lc_namespace': ['langc...</td>\n",
" <td>3.970739</td>\n",
" <td>176</td>\n",
" <td>0.000208</td>\n",
" <td>28.626591</td>\n",
" <td>15.4</td>\n",
" <td>18.9</td>\n",
" <td>5</td>\n",
" <td>The output provided by the model is a detailed...</td>\n",
" </tr>\n",
" <tr>\n",
" <th>1</th>\n",
" <td>How to run MLflow.evaluate()?</td>\n",
" <td>\\n\\nYou can run MLflow.evaluate() by using the...</td>\n",
" <td>How to run MLflow.evaluate()?</td>\n",
" <td>[{'lc_attributes': {}, 'lc_namespace': ['langc...</td>\n",
" <td>1.083653</td>\n",
" <td>39</td>\n",
" <td>0.000179</td>\n",
" <td>44.533493</td>\n",
" <td>4.7</td>\n",
" <td>4.5</td>\n",
" <td>5</td>\n",
" <td>The output states that \"You can run MLflow.eva...</td>\n",
" </tr>\n",
" <tr>\n",
" <th>2</th>\n",
" <td>How to log_table()?</td>\n",
" <td>\\n\\nYou can use the log_table() function in ML...</td>\n",
" <td>How to log_table()?</td>\n",
" <td>[{'lc_attributes': {}, 'lc_namespace': ['langc...</td>\n",
" <td>2.833117</td>\n",
" <td>114</td>\n",
" <td>0.000564</td>\n",
" <td>13.269521</td>\n",
" <td>7.9</td>\n",
" <td>8.8</td>\n",
" <td>1</td>\n",
" <td>The output provides a detailed explanation of ...</td>\n",
" </tr>\n",
" <tr>\n",
" <th>3</th>\n",
" <td>How to load_table()?</td>\n",
" <td>load_table() is not a function in MLflow.</td>\n",
" <td>How to load_table()?</td>\n",
" <td>[{'lc_attributes': {}, 'lc_namespace': ['langc...</td>\n",
" <td>3.736170</td>\n",
" <td>11</td>\n",
" <td>0.000144</td>\n",
" <td>193.916275</td>\n",
" <td>2.5</td>\n",
" <td>5.6</td>\n",
" <td>5</td>\n",
" <td>The output states that \"load_table() is not a ...</td>\n",
" </tr>\n",
" </tbody>\n",
"</table>\n",
"</div>"
],
"text/plain": [
" questions \\\n",
"0 What is MLflow? \n",
"1 How to run MLflow.evaluate()? \n",
"2 How to log_table()? \n",
"3 How to load_table()? \n",
"\n",
" outputs \\\n",
"0 MLflow is an open source platform for managin... \n",
"1 \\n\\nYou can run MLflow.evaluate() by using the... \n",
"2 \\n\\nYou can use the log_table() function in ML... \n",
"3 load_table() is not a function in MLflow. \n",
"\n",
" query \\\n",
"0 What is MLflow? \n",
"1 How to run MLflow.evaluate()? \n",
"2 How to log_table()? \n",
"3 How to load_table()? \n",
"\n",
" source_documents latency token_count \\\n",
"0 [{'lc_attributes': {}, 'lc_namespace': ['langc... 3.970739 176 \n",
"1 [{'lc_attributes': {}, 'lc_namespace': ['langc... 1.083653 39 \n",
"2 [{'lc_attributes': {}, 'lc_namespace': ['langc... 2.833117 114 \n",
"3 [{'lc_attributes': {}, 'lc_namespace': ['langc... 3.736170 11 \n",
"\n",
" toxicity/v1/score perplexity/v1/score \\\n",
"0 0.000208 28.626591 \n",
"1 0.000179 44.533493 \n",
"2 0.000564 13.269521 \n",
"3 0.000144 193.916275 \n",
"\n",
" flesch_kincaid_grade_level/v1/score ari_grade_level/v1/score \\\n",
"0 15.4 18.9 \n",
"1 4.7 4.5 \n",
"2 7.9 8.8 \n",
"3 2.5 5.6 \n",
"\n",
" faithfulness/v1/score faithfulness/v1/justification \n",
"0 5 The output provided by the model is a detailed... \n",
"1 5 The output states that \"You can run MLflow.eva... \n",
"2 1 The output provides a detailed explanation of ... \n",
"3 5 The output states that \"load_table() is not a ... "
]
},
"execution_count": 18,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"results.tables[\"eval_results_table\"]"
]
}
],
"metadata": {
"application/vnd.databricks.v1+notebook": {
"dashboards": [],
"language": "python",
"notebookMetadata": {
"pythonIndentUnit": 2
},
"notebookName": "LLM Evaluation Examples -- RAG",
"widgets": {}
},
"kernelspec": {
"display_name": "mlflow-dev-env",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.8.17"
}
},
"nbformat": 4,
"nbformat_minor": 0
}
+22
View File
@@ -0,0 +1,22 @@
name: FlowerClassifier
python_env: python_env.yaml
entry_points:
# train Keras DL model
main:
parameters:
training_data: {type: string, default: "./flower_photos"}
epochs: {type: int, default: 1}
image_width: {type: int, default: 224}
image_height: {type: int, default: 224}
batch_size: {type: int, default: 16}
test_ratio: {type: float, default: 0.2}
seed: {type: int, default: 97531}
command: "python train.py --training-data {training_data}
--batch-size {batch_size}
--epochs {epochs}
--image-width {image_width}
--image-height {image_height}
--test-ratio {test_ratio}"
+107
View File
@@ -0,0 +1,107 @@
How To Train and Deploy Image Classifier with MLflow and Keras
--------------------------------------------------------------
In this example we demonstrate how to train and deploy image classification models with MLflow.
We train a VGG16 deep learning model to classify flower species from photos using a `dataset
<http://download.tensorflow.org/example_images/flower_photos.tgz>`_ available from `tensorflow.org
<http://www.tensorflow.org>`_. Note that although we use Keras to train the model in this case,
a similar approach can be applied to other deep learning frameworks such as ``PyTorch``.
The MLflow model produced by running this example can be deployed to any MLflow supported endpoints.
All the necessary image preprocessing is packaged with the model. The model can therefore be applied
to image data directly. All that is required in order to pass new data to the model is to encode the
image binary data as base64 encoded string in pandas DataFrame (standard interface for MLflow python
function models). The included Python scripts demonstrate how the model can be deployed to a REST
API endpoint for realtime evaluation or to Spark for batch scoring..
In order to include custom image pre-processing logic with the model, we define the model as a
custom python function model wrapping around the underlying Keras model. The wrapper provides
necessary preprocessing to convert input data into multidimensional arrays expected by the
Keras model. The preprocessing logic is stored with the model as a code dependency. Here is an
example of the output model directory layout:
.. code-block:: bash
tree model
::
model
├── MLmodel
├── code
│   └── image_pyfunc.py
├── data
│   └── image_model
│   ├── conf.yaml
│   └── keras_model
│   ├── MLmodel
│   ├── conda.yaml
│   └── model.h5
└── mlflow_env.yml
The example contains the following files:
* MLproject
Contains definition of this project. Contains only one entry point to train the model.
* conda.yaml
Defines project dependencies. NOTE: You might want to change tensorflow package to tensorflow-gpu
if you have gpu(s) available.
* train.py
Main entry point of the projects. Handles command line arguments and possibly downloads the
dataset.
* image_pyfunc.py
The implementation of the model train and also of the outputed custom python flavor model. Note
that the same preprocessing code that is used during model training is packaged with the output
model and is used during scoring.
* score_images_rest.py
Score an image or a directory of images using a model deployed to a REST endpoint.
* score_images_spark.py
Score an image or a directory of images using model deployed to Spark.
Running this Example
^^^^^^^^^^^^^^^^^^^^
To train the model, run the example as a standard MLflow project:
.. code-block:: bash
mlflow run examples/flower_classifier
This will download the training dataset from ``tensorflow.org``, train a classifier using Keras and
log results with MLflow.
To test your model, run the included scoring scripts. For example, say your model was trained with
run_id ``101``.
- To test REST api scoring do the following two steps:
1. Deploy the model as a local REST endpoint by running ``mlflow models serve``:
.. code-block:: bash
# deploy the model to local REST api endpoint
mlflow models serve --model-uri runs:/101/model --port 54321
1. Apply the model to new data using the provided score_images_rest.py script:
.. code-block:: bash
# score the deployed model
python score_images_rest.py --host http://127.0.0.1 --port 54321 /path/to/images/for/scoring
- To test batch scoring in Spark, run score_images_spark.py to score the model in Spark like this:
.. code-block:: bash
python score_images_spark.py --model-uri runs:/101/model /path/to/images/for/scoring
+193
View File
@@ -0,0 +1,193 @@
"""
Example of a custom python function implementing image classifier with image preprocessing embedded
in the model.
"""
import base64
import importlib.metadata
import os
from io import BytesIO
from typing import Any
import keras
import numpy as np
import pandas as pd
import PIL
import tensorflow as tf
import yaml
from PIL import Image
import mlflow
from mlflow.utils import PYTHON_VERSION
from mlflow.utils.file_utils import TempDir
def decode_and_resize_image(raw_bytes, size):
"""
Read, decode and resize raw image bytes (e.g. raw content of a jpeg file).
Args:
raw_bytes: Image bits, e.g. jpeg image.
size: Requested output dimensions.
Returns:
Multidimensional numpy array representing the resized image.
"""
return np.asarray(Image.open(BytesIO(raw_bytes)).resize(size), dtype=np.float32)
class KerasImageClassifierPyfunc:
"""
Image classification model with embedded pre-processing.
This class is essentially an MLflow custom python function wrapper around a Keras model.
The wrapper provides image preprocessing so that the model can be applied to images directly.
The input to the model is base64 encoded image binary data (e.g. contents of a jpeg file).
The output is the predicted class label, predicted class id followed by probabilities for each
class.
The model declares current local versions of Keras, Tensorlow and pillow as dependencies in its
conda environment file.
"""
def __init__(self, graph, session, model, image_dims, domain):
self._graph = graph
self._session = session
self._model = model
self._image_dims = image_dims
self._domain = domain
probs_names = [f"p({x})" for x in domain]
self._column_names = ["predicted_label", "predicted_label_id"] + probs_names
def predict(
self,
input,
params: dict[str, Any] | None = None,
):
"""
Generate predictions for the data.
Args:
input: pandas.DataFrame with one column containing images to be scored. The image
column must contain base64 encoded binary content of the image files. The image
format must be supported by PIL (e.g. jpeg or png).
params: Additional parameters to pass to the model for inference.
Returns:
pandas.DataFrame containing predictions with the following schema:
Predicted class: string,
Predicted class index: int,
Probability(class==0): float,
...,
Probability(class==N): float,
"""
# decode image bytes from base64 encoding
def decode_img(x):
return pd.Series(base64.decodebytes(bytearray(x[0], encoding="utf8")))
images = input.apply(axis=1, func=decode_img)
probs = self._predict_images(images)
m, n = probs.shape
label_idx = np.argmax(probs, axis=1)
labels = np.array([self._domain[i] for i in label_idx], dtype=str).reshape(m, 1)
output_data = np.concatenate((labels, label_idx.reshape(m, 1), probs), axis=1)
res = pd.DataFrame(columns=self._column_names, data=output_data)
res.index = input.index
return res
def _predict_images(self, images):
"""
Generate predictions for input images.
Args:
images: Binary image data.
Returns:
Predicted probabilities for each class.
"""
def preprocess_f(z):
return decode_and_resize_image(z, self._image_dims[:2])
x = np.array(images[images.columns[0]].apply(preprocess_f).tolist())
with self._graph.as_default():
with self._session.as_default():
return self._model.predict(x)
def log_model(keras_model, signature, artifact_path, image_dims, domain):
"""
Log a KerasImageClassifierPyfunc model as an MLflow artifact for the current run.
Args:
keras_model: Keras model to be saved.
signature: Model signature.
artifact_path: Run-relative artifact path this model is to be saved to.
image_dims: Image dimensions the Keras model expects.
domain: Labels for the classes this model can predict.
"""
with TempDir() as tmp:
data_path = tmp.path("image_model")
os.mkdir(data_path)
conf = {"image_dims": "/".join(map(str, image_dims)), "domain": "/".join(map(str, domain))}
with open(os.path.join(data_path, "conf.yaml"), "w") as f:
yaml.safe_dump(conf, stream=f)
keras_path = os.path.join(data_path, "keras_model")
mlflow.tensorflow.save_model(model=keras_model, path=keras_path)
conda_env = tmp.path("conda_env.yaml")
with open(conda_env, "w") as f:
f.write(
conda_env_template.format(
python_version=PYTHON_VERSION,
keras_version=keras.__version__,
tf_name=tf.__name__, # can have optional -gpu suffix
tf_version=tf.__version__,
pip_version=importlib.metadata.version("pip"),
pillow_version=PIL.__version__,
)
)
mlflow.pyfunc.log_model(
name=artifact_path,
signature=signature,
loader_module=__name__,
code_paths=[__file__],
data_path=data_path,
conda_env=conda_env,
)
def _load_pyfunc(path):
"""
Load the KerasImageClassifierPyfunc model.
"""
with open(os.path.join(path, "conf.yaml")) as f:
conf = yaml.safe_load(f)
keras_model_path = os.path.join(path, "keras_model")
domain = conf["domain"].split("/")
image_dims = np.array([int(x) for x in conf["image_dims"].split("/")], dtype=np.int32)
# NOTE: TensorFlow based models depend on global state (Graph and Session) given by the context.
# To make sure we score the model in the same session as we loaded it in, we create a new
# session and a new graph here and store them with the model.
with tf.Graph().as_default() as g:
with tf.Session().as_default() as sess:
keras.backend.set_session(sess)
keras_model = mlflow.tensorflow.load_model(keras_model_path)
return KerasImageClassifierPyfunc(g, sess, keras_model, image_dims, domain=domain)
conda_env_template = """
name: flower_classifier
channels:
- conda-forge
dependencies:
- python=={python_version}
- pip=={pip_version}
- pip:
- mlflow>=1.6
- pillow=={pillow_version}
- keras=={keras_version}
- {tf_name}=={tf_version}
"""
@@ -0,0 +1,8 @@
build_dependencies:
- pip==22.2.2
dependencies:
- mlflow>=1.6
- pandas==1.5.0
- scikit-learn==1.1.3
- tensorflow==2.10.0
- pillow==9.2.0
@@ -0,0 +1,71 @@
"""
Example of scoring images with MLflow model deployed to a REST API endpoint.
The MLflow model to be scored is expected to be an instance of KerasImageClassifierPyfunc
(e.g. produced by running this project) and deployed with MLflow prior to invoking this script.
"""
import base64
import os
import click
import pandas as pd
import requests
from mlflow.utils import cli_args
def score_model(path, host, port):
"""
Score images on the local path with MLflow model deployed at given uri and port.
Args:
path: Path to a single image file or a directory of images.
host: Host the model is deployed at.
port: Port the model is deployed at.
Returns:
Server response.
"""
if os.path.isdir(path):
filenames = [
os.path.join(path, x) for x in os.listdir(path) if os.path.isfile(os.path.join(path, x))
]
else:
filenames = [path]
def read_image(x):
with open(x, "rb") as f:
return f.read()
data = pd.DataFrame(
data=[base64.encodebytes(read_image(x)) for x in filenames], columns=["image"]
).to_json(orient="split")
response = requests.post(
url=f"{host}:{port}/invocations",
data={
"dataframe_split": data,
},
headers={"Content-Type": "application/json"},
)
if response.status_code != 200:
raise Exception(f"Status Code {response.status_code}. {response.text}")
return response
@click.command(help="Score images.")
@click.option("--port", type=click.INT, default=80, help="Port at which the model is deployed.")
@cli_args.HOST
@click.argument("data-path")
def run(data_path, host, port):
"""
Score images with MLflow deployed deployed at given uri and port and print out the response
to standard out.
"""
print(score_model(data_path, host, port).text)
if __name__ == "__main__":
run()
@@ -0,0 +1,97 @@
"""
Example of scoring images with MLflow model produced by running this project in Spark.
The MLflow model is loaded to Spark using ``mlflow.pyfunc.spark_udf``. The images are read as binary
data and represented as base64 encoded string column and passed to the model. The results are
returned as a column with predicted class label, class id and probabilities for each class encoded
as an array of strings.
"""
import base64
import os
import click
import pandas as pd
import pyspark
from pyspark.sql.types import ArrayType, Row, StringType, StructField, StructType
import mlflow.pyfunc
from mlflow.utils import cli_args
def read_image_bytes_base64(path):
with open(path, "rb") as f:
return str(base64.encodebytes(f.read()), encoding="utf8")
def read_images(spark, filenames):
filenames_rdd = spark.sparkContext.parallelize(filenames)
schema = StructType([
StructField("filename", StringType(), True),
StructField("image", StringType(), True),
])
return filenames_rdd.map(lambda x: Row(filename=x, image=read_image_bytes_base64(x))).toDF(
schema=schema
)
def score_model(spark, data_path, model_uri):
if os.path.isdir(data_path):
filenames = [
os.path.abspath(os.path.join(data_path, x))
for x in os.listdir(data_path)
if os.path.isfile(os.path.join(data_path, x))
]
else:
filenames = [data_path]
image_classifier_udf = mlflow.pyfunc.spark_udf(
spark=spark, model_uri=model_uri, result_type=ArrayType(StringType())
)
image_df = read_images(spark, filenames)
raw_preds = (
image_df
.withColumn("prediction", image_classifier_udf("image"))
.select(["filename", "prediction"])
.toPandas()
)
# load the pyfunc model to get our domain
pyfunc_model = mlflow.pyfunc.load_model(model_uri=model_uri)
preds = pd.DataFrame(raw_preds["filename"], index=raw_preds.index)
preds[pyfunc_model._column_names] = pd.DataFrame(
raw_preds["prediction"].values.tolist(),
columns=pyfunc_model._column_names,
index=raw_preds.index,
)
preds = pd.DataFrame(raw_preds["filename"], index=raw_preds.index)
preds[pyfunc_model._column_names] = pd.DataFrame(
raw_preds["prediction"].values.tolist(),
columns=pyfunc_model._column_names,
index=raw_preds.index,
)
return preds.to_json(orient="records")
@click.command(help="Score images.")
@cli_args.MODEL_URI
@click.argument("data-path")
def run(data_path, model_uri):
with (
pyspark.sql.SparkSession.builder
.config(key="spark.python.worker.reuse", value=True)
.config(key="spark.ui.enabled", value=False)
.master("local-cluster[2, 1, 1024]")
.getOrCreate() as spark
):
# ignore spark log output
spark.sparkContext.setLogLevel("OFF")
print(score_model(spark, data_path, model_uri))
if __name__ == "__main__":
run()
+244
View File
@@ -0,0 +1,244 @@
"""
Example of image classification with MLflow using Keras to classify flowers from photos. The data is
taken from ``http://download.tensorflow.org/example_images/flower_photos.tgz`` and may be
downloaded during running this project if it is missing.
"""
import math
import os
import tarfile
import click
import keras
import numpy as np
import tensorflow as tf
from image_pyfunc import decode_and_resize_image, log_model
from keras.applications import vgg16
from keras.callbacks import Callback
from keras.layers import Dense, Flatten, Input, Lambda
from keras.models import Model
from keras.utils import np_utils
from sklearn.model_selection import train_test_split
import mlflow
from mlflow.models import infer_signature
def download_input():
import requests
url = "http://download.tensorflow.org/example_images/flower_photos.tgz"
print("downloading '{}' into '{}'".format(url, os.path.abspath("flower_photos.tgz")))
r = requests.get(url)
with open("flower_photos.tgz", "wb") as f:
f.write(r.content)
print("decompressing flower_photos.tgz to '{}'".format(os.path.abspath("flower_photos")))
with tarfile.open("flower_photos.tgz") as tar:
tar.extractall(path="./")
@click.command(
help="Trains an Keras model on flower_photos dataset. "
"The input is expected as a directory tree with pictures for each category in a "
"folder named by the category. "
"The model and its metrics are logged with mlflow."
)
@click.option("--epochs", type=click.INT, default=1, help="Maximum number of epochs to evaluate.")
@click.option(
"--batch-size", type=click.INT, default=16, help="Batch size passed to the learning algo."
)
@click.option("--image-width", type=click.INT, default=224, help="Input image width in pixels.")
@click.option("--image-height", type=click.INT, default=224, help="Input image height in pixels.")
@click.option("--seed", type=click.INT, default=97531, help="Seed for the random generator.")
@click.option("--training-data", type=click.STRING, default="./flower_photos")
@click.option("--test-ratio", type=click.FLOAT, default=0.2)
def run(training_data, test_ratio, epochs, batch_size, image_width, image_height, seed):
image_files = []
labels = []
domain = {}
print("Training model with the following parameters:")
for param, value in locals().items():
print(" ", param, "=", value)
if training_data == "./flower_photos" and not os.path.exists(training_data):
print("Input data not found, attempting to download the data from the web.")
download_input()
for dirname, _, files in os.walk(training_data):
for filename in files:
if filename.endswith("jpg"):
image_files.append(os.path.join(dirname, filename))
clazz = os.path.basename(dirname)
if clazz not in domain:
domain[clazz] = len(domain)
labels.append(domain[clazz])
train(
image_files,
labels,
domain,
epochs=epochs,
test_ratio=test_ratio,
batch_size=batch_size,
image_width=image_width,
image_height=image_height,
seed=seed,
)
class MlflowLogger(Callback):
"""
Keras callback for logging metrics and final model with MLflow.
Metrics are logged after every epoch. The logger keeps track of the best model based on the
validation metric. At the end of the training, the best model is logged with MLflow.
"""
def __init__(self, model, x_train, y_train, x_valid, y_valid, **kwargs):
self._model = model
self._best_val_loss = math.inf
self._train = (x_train, y_train)
self._valid = (x_valid, y_valid)
self._pyfunc_params = kwargs
self._best_weights = None
def on_epoch_end(self, epoch, logs=None):
"""
Log Keras metrics with MLflow. Update the best model if the model improved on the validation
data.
"""
if not logs:
return
for name, value in logs.items():
name = "valid_" + name[4:] if name.startswith("val_") else "train_" + name
mlflow.log_metric(name, value)
val_loss = logs["val_loss"]
if val_loss < self._best_val_loss:
# Save the "best" weights
self._best_val_loss = val_loss
self._best_weights = [x.copy() for x in self._model.get_weights()]
def on_train_end(self, *args, **kwargs):
"""
Log the best model with MLflow and evaluate it on the train and validation data so that the
metrics stored with MLflow reflect the logged model.
"""
self._model.set_weights(self._best_weights)
x, y = self._train
train_res = self._model.evaluate(x=x, y=y)
for name, value in zip(self._model.metrics_names, train_res):
mlflow.log_metric(f"train_{name}", value)
x, y = self._valid
valid_res = self._model.evaluate(x=x, y=y)
for name, value in zip(self._model.metrics_names, valid_res):
mlflow.log_metric(f"valid_{name}", value)
signature = infer_signature(x, y)
log_model(keras_model=self._model, signature=signature, **self._pyfunc_params)
def _imagenet_preprocess_tf(x):
return (x / 127.5) - 1
def _create_model(input_shape, classes):
image = Input(input_shape)
lambda_layer = Lambda(_imagenet_preprocess_tf)
preprocessed_image = lambda_layer(image)
model = vgg16.VGG16(
classes=classes, input_tensor=preprocessed_image, weights=None, include_top=False
)
x = Flatten(name="flatten")(model.output)
x = Dense(4096, activation="relu", name="fc1")(x)
x = Dense(4096, activation="relu", name="fc2")(x)
x = Dense(classes, activation="softmax", name="predictions")(x)
return Model(inputs=model.input, outputs=x)
def train(
image_files,
labels,
domain,
image_width=224,
image_height=224,
epochs=1,
batch_size=16,
test_ratio=0.2,
seed=None,
):
"""
Train VGG16 model on provided image files. This will create a new MLflow run and log all
parameters, metrics and the resulting model with MLflow. The resulting model is an instance
of KerasImageClassifierPyfunc - a custom python function model that embeds all necessary
preprocessing together with the VGG16 Keras model. The resulting model can be applied
directly to image base64 encoded image data.
Args:
image_files: List of image files to be used for training.
labels: List of labels for the image files.
domain: Dictionary representing the domain of the response.
Provides mapping label-name -> label-id.
image_width: Width of the input image in pixels.
image_height: Height of the input image in pixels.
epochs: Number of epochs to train the model for.
batch_size: Batch size used during training.
test_ratio: Fraction of dataset to be used for validation. This data will not be used
during training.
seed: Random seed. Used e.g. when splitting the dataset into train / validation.
"""
assert len(set(labels)) == len(domain)
input_shape = (image_width, image_height, 3)
with mlflow.start_run():
mlflow.log_param("epochs", str(epochs))
mlflow.log_param("batch_size", str(batch_size))
mlflow.log_param("validation_ratio", str(test_ratio))
if seed:
mlflow.log_param("seed", str(seed))
def _read_image(filename):
with open(filename, "rb") as f:
return f.read()
with tf.Graph().as_default() as g:
with tf.compat.v1.Session(graph=g).as_default():
dims = input_shape[:2]
x = np.array([decode_and_resize_image(_read_image(x), dims) for x in image_files])
y = np_utils.to_categorical(np.array(labels), num_classes=len(domain))
train_size = 1 - test_ratio
x_train, x_valid, y_train, y_valid = train_test_split(
x, y, random_state=seed, train_size=train_size
)
model = _create_model(input_shape=input_shape, classes=len(domain))
model.compile(
optimizer=keras.optimizers.SGD(decay=1e-5, nesterov=True, momentum=0.9),
loss=keras.losses.categorical_crossentropy,
metrics=["accuracy"],
)
sorted_domain = sorted(domain.keys(), key=lambda x: domain[x])
model.fit(
x=x_train,
y=y_train,
validation_data=(x_valid, y_valid),
epochs=epochs,
batch_size=batch_size,
callbacks=[
MlflowLogger(
model=model,
x_train=x_train,
y_train=y_train,
x_valid=x_valid,
y_valid=y_valid,
artifact_path="model",
domain=sorted_domain,
image_dims=input_shape,
)
],
)
if __name__ == "__main__":
run()
+82
View File
@@ -0,0 +1,82 @@
# MLflow AI Gateway
The examples provided within this directory show how to get started with individual providers and at least
one of the supported endpoint types. When configuring an instance of the MLflow AI Gateway, multiple providers,
instances of endpoint types, and model versions can be specified for each query endpoint on the server.
## Example configuration files
Within this directory are example config files for each of the supported providers. If using these as a guide
for configuring a large number of endpoints, ensure that the placeholder names (i.e., "completions", "chat", "embeddings")
are modified to prevent collisions. These names are provided for clarity only for the examples and real-world
use cases should define a relevant and meaningful endpoint name to eliminate ambiguity and minimize the chances of name collisions.
# Getting Started with MLflow AI Gateway for OpenAI
This guide will walk you through the installation and basic setup of the MLflow AI Gateway.
Within sub directories of this examples section, you can find specific executable examples
that can be used to validate a given provider's configuration through the MLflow AI Gateway.
Let's get started.
## Step 1: Installing the MLflow AI Gateway
The MLflow AI Gateway is best installed from PyPI. Open your terminal and use the following pip command:
```sh
# Installation from PyPI
pip install 'mlflow[genai]'
```
For those interested in development or in using the most recent build of the MLflow AI Gateway, you may choose to install from the fork of the repository:
```sh
# Installation from the repository
pip install -e '.[genai]'
```
## Step 2: Configuring Endpoints
Each provider has a distinct set of allowable endpoint types (i.e., chat, completions, etc) and
specific requirements for the initialization of the endpoints to interface with their services.
For full examples of configurations and supported endpoint types, see:
- [OpenAI](openai/config.yaml)
- [MosaicML](mosaicml/config.yaml)
- [Anthropic](anthropic/config.yaml)
- [Cohere](cohere/config.yaml)
- [AI21 Labs](ai21labs/config.yaml)
- [PaLM](palm/config.yaml)
- [AzureOpenAI](azure_openai/config.yaml)
- [Mistral](mistral/config.yaml)
- [TogetherAI](togetherai/config.yaml)
## Step 3: Setting Access Keys
See information on specific methods of obtaining and setting the access keys within the provider-specific documentation within this directory.
## Step 4: Starting the MLflow AI Gateway
With the MLflow configuration file in place and access key(s) set, you can now start the MLflow AI Gateway.
Replace `<provider>` with the actual path to the MLflow configuration file for the provider of your choice:
```sh
mlflow gateway start --config-path examples/gateway/<provider>/config.yaml --port 7000
# For example:
mlflow gateway start --config-path examples/gateway/openai/config.yaml --port 7000
```
## Step 5: Accessing the Interactive API Documentation
With the MLflow AI Gateway up and running, access its interactive API documentation by navigating to the following URL:
http://127.0.0.1:7000/docs
## Step 6: Sending Test Requests
After successfully setting up the MLflow AI Gateway, you can send a test request using the provided Python script.
Replace <provider> with the name of the provider example test script that you'd like to use:
```sh
python examples/gateway/<provider>/example.py
```
+13
View File
@@ -0,0 +1,13 @@
## Example endpoint configuration for AI21 Labs
To set up your MLflow configuration file, include a single endpoint for the completions endpoint as shown in the [AI21 labs configuration](config.yaml) YAML file.
## Obtaining and Setting the AI21 Labs API Key
To obtain an AI21 Labs API key, you need to create an account and subscribe to the service at [AI21 Labs](https://studio.ai21.com/account/api-key?source=docs).
After obtaining the key, you can export it to your environment variables. Make sure to replace the '...' with your actual API key:
```sh
export AI21LABS_API_KEY=...
```
+8
View File
@@ -0,0 +1,8 @@
endpoints:
- name: completions
endpoint_type: llm/v1/completions
model:
provider: ai21labs
name: j2-mid
config:
ai21labs_api_key: $AI21LABS_API_KEY
+22
View File
@@ -0,0 +1,22 @@
from mlflow.deployments import get_deploy_client
def main():
client = get_deploy_client("http://localhost:7000")
print(f"AI21 Labs endpoints: {client.list_endpoints()}\n")
print(f"AI21 Labs completions endpoint info: {client.get_endpoint(endpoint='completions')}\n")
# Completions request
response_completions = client.predict(
endpoint="completions",
inputs={
"prompt": "What is the world record for flapjack consumption in a single sitting?",
"temperature": 0.1,
},
)
print(f"AI21 Labs response for completions: {response_completions}")
if __name__ == "__main__":
main()
+13
View File
@@ -0,0 +1,13 @@
## Example endpoint configuration for Anthropic
To set up your MLflow configuration file, include a single endpoint for the completions endpoint as shown in the [anthropic configuration](config.yaml) YAML file.
## Obtaining and Setting the Anthropic API Key
To obtain an Anthropic API key, you need to create an account and subscribe to the service at [Anthropic](https://docs.anthropic.com/claude/docs/getting-access-to-claude).
After obtaining the key, you can export it to your environment variables. Make sure to replace the '...' with your actual API key:
```sh
export ANTHROPIC_API_KEY=...
```
+8
View File
@@ -0,0 +1,8 @@
endpoints:
- name: completions
endpoint_type: llm/v1/completions
model:
provider: anthropic
name: claude-1.3-100k
config:
anthropic_api_key: $ANTHROPIC_API_KEY
+23
View File
@@ -0,0 +1,23 @@
from mlflow.deployments import get_deploy_client
def main():
client = get_deploy_client("http://localhost:7000")
print(f"Anthropic endpoints: {client.list_endpoints()}\n")
print(f"Anthropic completions endpoint info: {client.get_endpoint(endpoint='completions')}\n")
# Completions request
response_completions = client.predict(
endpoint="completions",
inputs={
"prompt": "How many average size European ferrets can fit inside a standard olympic "
"size swimming pool?",
"max_tokens": 5000,
},
)
print(f"Anthropic response for completions: {response_completions}")
if __name__ == "__main__":
main()
+26
View File
@@ -0,0 +1,26 @@
## Example endpoint configuration for Azure OpenAI
The following example configuration shows the 3 supported endpoints for Azure OpenAI: chat, completions, and embeddings.
Additionally, it illustrates the two separate api types that are supported for this service.
- `azure` api type: uses a generated token that is applied by setting the API token key directly to an environment variable
- `azuread` api type: uses Azure Active Directory for supplying the active directory key to be used to an environment variable
Depending on how your users will be interacting with the MLflow AI Gateway, a single access paradigm (either `azure` **or** `azuread` is recommended, not a mix of both).
See the [Azure OpenAI configuration](config.yaml) YAML file for example configurations showing all supported endpoint types and the different token access types.
## Setting the Azure OpenAI API Key
In order to get access to the Azure OpenAI service, [see the documentation](https://azure.microsoft.com/en-us/products/cognitive-services/openai-service) guidance in the cognitive services portal.
With the key, export it to your environment variables.
Replace the '...' with your actual API key:
```sh
export OPENAI_API_KEY=...
```
## Validating the Azure OpenAI endpoint
See the [OpenAI Example](../openai/example.py) for testing the Azure OpenAI endpoints. The usage is identical to the standard OpenAI integration from an API perspective.
+36
View File
@@ -0,0 +1,36 @@
endpoints:
- name: chat
endpoint_type: llm/v1/chat
model:
provider: openai
name: gpt-4o-mini
config:
openai_api_type: "azure"
openai_api_key: $OPENAI_API_KEY
openai_deployment_name: "{your_deployment_name}"
openai_api_base: "https://{your_resource_name}-azureopenai.openai.azure.com/"
openai_api_version: "2023-05-15"
- name: completions
endpoint_type: llm/v1/completions
model:
provider: openai
name: gpt-4o-mini
config:
openai_api_type: "azuread"
openai_api_key: $AZURE_AAD_TOKEN
openai_deployment_name: "{your_deployment_name}"
openai_api_base: "https://{your_resource_name}-azureopenai.openai.azure.com/"
openai_api_version: "2023-05-15"
- name: embeddings
endpoint_type: llm/v1/embeddings
model:
provider: openai
name: text-embedding-ada-002
config:
openai_api_type: "azure"
openai_api_key: $OPENAI_API_KEY
openai_deployment_name: "{your_deployment_name}"
openai_api_base: "https://{your_resource_name}-azureopenai.openai.azure.com/"
openai_api_version: "2023-05-15"
+7
View File
@@ -0,0 +1,7 @@
## Example endpoint configuration for Amazon Bedrock
To view an example of a Bedrock endpoint configuration, see [the configuration example](config.yaml) YAML file.
## Credentials
Valid AWS credentials are required for this example. Set `AWS_ACCESS_KEY_ID` and `AWS_SECRET_ACCESS_KEY` to valid credentials, or run in an environment with those variables set.
+11
View File
@@ -0,0 +1,11 @@
endpoints:
- name: completions
endpoint_type: llm/v1/completions
model:
provider: amazon-bedrock
name: amazon.titan-tg1-large
config:
aws_config:
aws_region: us-east-1
aws_access_key_id: $AWS_ACCESS_KEY_ID
aws_secret_access_key: $AWS_SECRET_ACCESS_KEY
+23
View File
@@ -0,0 +1,23 @@
from mlflow.deployments import get_deploy_client
def main():
client = get_deploy_client("http://localhost:7000")
print(f"Bedrock endpoints: {client.list_endpoints()}\n")
print(f"Bedrock completions endpoint info: {client.get_endpoint(endpoint='completions')}\n")
# Completions example
response_completions = client.predict(
endpoint="completions",
inputs={
"prompt": "How many patties could be stacked on a cheeseburger before issues arise?",
"max_tokens": 200,
"temperature": 0.25,
},
)
print(f"Bedrock completions response: {response_completions}")
if __name__ == "__main__":
main()
+13
View File
@@ -0,0 +1,13 @@
## Example endpoint configuration for Cohere
To see an example of specifying both the completions and the embeddings endpoints for Cohere, see [the configuration](config.yaml) YAML file.
This configuration file specifies two endpoints: 'completions' and 'embeddings', both using Cohere's models 'command' and 'embed-english-light-v2.0', respectively.
## Setting a Cohere API Key
This example requires a [Cohere API key](https://docs.cohere.com/docs/going-live):
```sh
export COHERE_API_KEY=...
```
+16
View File
@@ -0,0 +1,16 @@
endpoints:
- name: completions
endpoint_type: llm/v1/completions
model:
provider: cohere
name: command
config:
cohere_api_key: $COHERE_API_KEY
- name: embeddings
endpoint_type: llm/v1/embeddings
model:
provider: cohere
name: embed-english-light-v2.0
config:
cohere_api_key: $COHERE_API_KEY
+29
View File
@@ -0,0 +1,29 @@
from mlflow.deployments import get_deploy_client
def main():
client = get_deploy_client("http://localhost:7000")
print(f"Cohere endpoints: {client.list_endpoints()}\n")
print(f"Cohere completions endpoint info: {client.get_endpoint(endpoint='completions')}\n")
# Completions request
response_completions = client.predict(
endpoint="completions",
inputs={
"prompt": "What is the world record for flapjack consumption in a single sitting?",
"temperature": 0.1,
},
)
print(f"Cohere response for completions: {response_completions}")
# Embeddings request
response_embeddings = client.predict(
endpoint="embeddings",
inputs={"input": ["Do you carry the Storm Trooper costume in size 2T?"]},
)
print(f"Cohere response for embeddings: {response_embeddings}")
if __name__ == "__main__":
main()
+13
View File
@@ -0,0 +1,13 @@
## Example endpoint configuration for GEMINI
To see an example of specifying both the completions and embeddings endpoints for Gemini, see [the configuration](config.yaml) YAML file.
This configuration file specifies three endpoints: 'completions', 'embeddings', and 'chat', using Gemini's model gemini-2.0-flash for completions and chat and gemini-embedding-exp-03-07 for embeddings.
## Setting a GEMINI API Key
This example requires a [GEMINI API key](https://ai.google.dev/gemini-api/docs/api-key):
```sh
export GEMINI_API_KEY=...
```
+24
View File
@@ -0,0 +1,24 @@
endpoints:
- name: embeddings
endpoint_type: llm/v1/embeddings
model:
provider: gemini
name: gemini-embedding-exp-03-07
config:
gemini_api_key: $GEMINI_API_KEY
- name: completions
endpoint_type: llm/v1/completions
model:
provider: gemini
name: gemini-2.0-flash
config:
gemini_api_key: $GEMINI_API_KEY
- name: chat
endpoint_type: llm/v1/chat
model:
provider: gemini
name: gemini-2.0-flash
config:
gemini_api_key: $GEMINI_API_KEY
+62
View File
@@ -0,0 +1,62 @@
from mlflow.deployments import get_deploy_client
def main():
client = get_deploy_client("http://localhost:7000")
print(f"Gemini endpoints: {client.list_endpoints()}\n")
print(f"Gemini completions endpoint info: {client.get_endpoint(endpoint='completions')}\n")
# Chat example
response_chat = client.predict(
endpoint="chat",
inputs={
"messages": [
{
"role": "system",
"content": "You are a talented European rapper with a background in US history",
},
{
"role": "user",
"content": "Please recite the preamble to the US Constitution as if it were "
"written today by a rapper from Reykjavík",
},
],
"temperature": 0.1,
"top_p": 1,
"n": 3,
"max_tokens": 1000,
"top_k": 40,
},
)
print(f"Gemini response for chat: {response_chat}")
# Embeddings request
response_embeddings = client.predict(
endpoint="embeddings",
inputs={
"input": [
"Describe the main differences between renewable and nonrenewable energy sources."
]
},
)
print(f"Gemini response for embeddings: {response_embeddings}\n")
# Completions request
response_completions = client.predict(
endpoint="completions",
inputs={
"prompt": "Describe the main differences between renewable and nonrenewable energy sources.",
"temperature": 0.1,
"stop": ["."],
"n": 3,
"max_tokens": 100,
"top_k": 40,
"top_p": 0.5,
},
)
print(f"Gemini response for completions: {response_completions}")
if __name__ == "__main__":
main()
+100
View File
@@ -0,0 +1,100 @@
## Example endpoint configuration for Huggingface Text Generation Inference
[Huggingface Text Generation Inference (TGI)](https://huggingface.co/docs/text-generation-inference/index) is a comprehensive toolkit designed for deploying and serving Large Language Models (LLMs) efficiently. It offers optimized support for various popular open-source LLMs such as Llama, Falcon, StarCoder, BLOOM, and GPT-Neo. TGI comes with various built-in optimizations and features, such as:
- Simple launcher to serve most popular LLMs
- Tensor Parallelism for faster inference on multiple GPUs
- Safetensors weight loading
- Optimized transformers code for inference using Flash Attention and Paged Attention on the most popular architectures
It should be noted that only a [selection of models](https://huggingface.co/docs/text-generation-inference/supported_models) are optimized for TGI, which uses custom CUDA kernels for faster inference. You can add the flag `--disable-custom-kernels`` at the end of the docker run command if you wish to disable them. If the above list lacks the model you would like to serve, or in the case you created a custom created model, you can try to initialize and serve the model anyways. However, since the model is not optimized for TGI, performance is not guaranteed.
For a more detailed description of all features, please go to the [documentation](https://huggingface.co/docs/text-generation-inference/index).
## Getting Started
> **NOTE** This example is tested on a Linux Machine (Debian 11) with a NVIDIA A100 GPU.
To configure the MLflow AI Gateway with Huggingface Text Generation Inference, a few additional steps need to be followed. The initial step involves deploying a Huggingface model on the TGI server, which is illustrated in the next section.
The recommended approach for deploying the TGI server is by utilizing the [official Docker container](ghcr.io/huggingface/text-generation-inference:1.1.1). Docker is an open-source platform that provides a streamlined solution for automating the deployment, scaling, and management of applications through containers. These containers encompass all the essential dependencies required for seamless execution, including libraries, binaries, and configuration files. To install Docker, please refer to the [installation guide](https://docs.docker.com/get-docker/).
Before proceeding, it is important to verify that your machine has the appropriate hardware to initiate the server. TGI optimized models are compatible with NVIDIA A100, A10G, and T4 GPUs. While other GPU hardware may still provide performance advantages, certain operations such as flash attention and paged attention will not be executed. If you intend to run the container on a machine lacking GPUs or CUDA support, you can eliminate the `--gpus all` flag and include `--disable-custom-kernels`. However, please note that the CPU is not the intended platform for the server, and this choice significantly impacts performance.
#### Installing the NVIDIA Container Toolkit
To begin, the installation of the NVIDIA container toolkit is necessary. This toolkit is essential for running GPU-accelerated containers. Execute the following command to acquire all the requisite packages [ref the code]:
```sh
curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | sudo gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg \
&& curl -s -L https://nvidia.github.io/libnvidia-container/stable/deb/nvidia-container-toolkit.list | \
sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' | \
sudo tee /etc/apt/sources.list.d/nvidia-container-toolkit.list \
&& \
sudo apt-get update
```
Install the NVIDIA Container toolkit by running the following command.
```
sudo apt-get install -y nvidia-container-toolkit
```
#### Running the TGI server.
After you installed the NVIDIA Container toolkit, you can run the following Docker command to to start a TGI server on your local machine on port `8000`. This will load a [falcon-7b-instruct](https://huggingface.co/tiiuae/falcon-7b-instruct) model on the TGI server.
```
model=tiiuae/falcon-7b-instruct
volume=$PWD/data # share a volume with the Docker container to avoid downloading weights every run
docker run --gpus all --shm-size 1g -p 8000:80 -v $volume:/data ghcr.io/huggingface/text-generation-inference:1.1.1 --model-id $model
```
After the TGI server is deployed, run the following script to verify that it is working correctly:
```
import requests
headers = {
"Content-Type": "application/json",
}
data = {
'inputs': 'What is Deep Learning?',
'parameters': {
'max_new_tokens': 20,
},
}
response = requests.post('http://127.0.0.1:8000/generate', headers=headers, json=data)
print(response.json())
# {'generated_text': '\nDeep learning is a branch of machine learning that uses artificial neural networks to learn and make decisions.'}
```
## Update the config.yaml to add a new embeddings endpoint
After you started the server, update the MLflow AI Gateway configuration file [config.yaml](config.yaml) and add the server as a new endpoint:
```
endpoints:
- name: completions
endpoint_type: llm/v1/completions
model:
provider: "huggingface-text-generation-inference"
name: llm
config:
hf_server_url: http://127.0.0.1:8000/generate
```
## Starting the MLflow AI Gateway
After the configuration file is created, you can start the MLflow AI Gateway by running the following command:
```
mlflow gateway start --config-path examples/gateway/huggingface/config.yaml --port 7000
```
## Querying the endpoint
See the [example script](example.py) within this directory to see how to query the `falcon-7b-instruct` model that is served.
## Setting the parameters of TGI
When you make a request to the MLflow Deployments server, the information you provide in the request body will be sent to TGI. This gives you more control over the output you receive from TGI. However, it's important to note that you cannot turn off `details` and `decoder_input_details`, as they are necessary for TGI endpoints to work correctly.
+8
View File
@@ -0,0 +1,8 @@
endpoints:
- name: completions
endpoint_type: llm/v1/completions
model:
provider: "huggingface-text-generation-inference"
name: falcon-7b-instruct
config:
hf_server_url: http://127.0.0.1:8080
+25
View File
@@ -0,0 +1,25 @@
from mlflow.deployments import get_deploy_client
def main():
client = get_deploy_client("http://localhost:7000")
print(f"Hugging Face TGI endpoints: {client.list_endpoints()}\n")
print(
f"Hugging Face completions endpoint info: {client.get_endpoint(endpoint='completions')}\n"
)
# Completions request
response_completions = client.predict(
endpoint="completions",
inputs={
"prompt": ("What is Deep Learning?"),
"temperature": 0.1,
},
)
print(f"Hugging Face TGI response for completions: {response_completions}")
if __name__ == "__main__":
main()
+13
View File
@@ -0,0 +1,13 @@
## Example endpoint configuration for Mistral
To see an example of specifying both the completions and the embeddings endpoints for Mistral, see [the configuration](config.yaml) YAML file.
This configuration file specifies two endpoints: 'completions' and 'embeddings', both using Mistral's models 'mistral-tiny' and 'mistral-embed', respectively.
## Setting a Mistral API Key
This example requires a [Mistral API key](https://docs.mistral.ai/):
```sh
export MISTRAL_API_KEY=...
```
+16
View File
@@ -0,0 +1,16 @@
endpoints:
- name: completions
endpoint_type: llm/v1/completions
model:
provider: mistral
name: mistral-tiny
config:
mistral_api_key: $MISTRAL_API_KEY
- name: embeddings
endpoint_type: llm/v1/embeddings
model:
provider: mistral
name: mistral-embed
config:
mistral_api_key: $MISTRAL_API_KEY
+34
View File
@@ -0,0 +1,34 @@
from mlflow.deployments import get_deploy_client
def main():
client = get_deploy_client("http://localhost:7000")
print(f"Mistral endpoints: {client.list_endpoints()}\n")
print(f"Mistral completions endpoint info: {client.get_endpoint(endpoint='completions')}\n")
# Completions request
response_completions = client.predict(
endpoint="completions",
inputs={
"prompt": "How many average size European ferrets can fit inside a standard olympic?",
"temperature": 0.1,
},
)
print(f"Mistral response for completions: {response_completions}")
# Embeddings request
response_embeddings = client.predict(
endpoint="embeddings",
inputs={
"input": [
"How does your culture celebrate the New Year, and how does it differ from other countries' "
"celebrations?"
]
},
)
print(f"Mistral response for embeddings: {response_embeddings}")
if __name__ == "__main__":
main()
+390
View File
@@ -0,0 +1,390 @@
# Guide to using an MLflow served model with MLflow Deployments
In order to utilize MLflow Deployments with MLflow model serving, a few steps must be taken
in addition to those for configuring access to SaaS models (such as Anthropic and OpenAI). The first and most obvious
step that must be taken prior to interfacing with an MLflow served model is that a model needs to be logged to the
MLflow tracking server.
An important consideration for deciding whether to interface MLflow Deployments with a specific model is to evaluate the PyFunc interface that the model will
return after being called for inference. Due to the fact that the MLflow AI Gateway defines a specific response signature, expectations for each endpoint type's payload contents
must be met in order for a endpoint to be valid.
For example, an embeddings endpoint (llm/v1/embeddings endpoint type) is designed to return embeddings data as a collection (a list) of floats that correspond to each of the
input strings that are sent for embeddings inference to a service. The expectation that the embeddings endpoint definition has is that the data is in a particular format. Specifically one that
is capable of having the embeddings data extractable from a service response. Therefore, an MLflow model that returns data in the format below is perfectly valid.
```json
{
"predictions": [
[0.0, 0.1],
[1.0, 0.0]
]
}
```
However, a return value from a serving endpoint via a custom PyFunc of the form below will not work.
```json
{
"predictions": [
{
"embedding": [0.0, 0.1]
},
{
"embedding": [1.0, 0.0]
}
]
}
```
It is important to note that the MLflow AI Gateway does not perform validation on a configured endpoint until the point of querying. Creating a endpoint that interfaces with the
MLflow model server that is returning a payload that is incompatible with the configured endpoint type definition will raise 502 exceptions only when queried.
> **NOTE:** It is important to validate the output response of a model served by MLflow to ensure compatibility with the MLflow Deployments endpoint definitions. Not all model outputs are compatible with given endpoint types.
## Creating and logging an embeddings model
To start, we need a model that is capable of generating embeddings. For this example, we'll use
the `sentence_transformers` library and the corresponding MLflow flavor.
```python
from sentence_transformers import SentenceTransformer
import mlflow
model = SentenceTransformer(model_name_or_path="all-MiniLM-L6-v2")
artifact_path = "embeddings_model"
with mlflow.start_run():
model_info = mlflow.sentence_transformers.log_model(
model,
name=artifact_path,
)
```
## Generate the cli command for starting a local MLflow Model Serving endpoint for this embeddings model
```python
print(f"mlflow models serve -m {model_info.model_uri} -h 127.0.0.1 -p 9020 --no-conda")
```
Copy the output from the print statement to the clipboard.
## Starting the model server for the embeddings model
With the printed string from running the above command copied to the clipboard, open a new terminal
and paste the string. Leave the terminal window open and running.
```commandline
mlflow models serve -m file:///Users/me/demos/mlruns/0/2bfcdcb66eaf4c88abe8e0c7bcab639e/artifacts/embeddings_model -h 127.0.0.1 -p 9020 --no-conda
```
## Update the config.yaml to add a new embeddings endpoint
After assigning a valid port and ensuring that the model server starts correctly:
```commandline
2023/08/08 17:36:44 INFO mlflow.models.flavor_backend_registry: Selected backend for flavor 'python_function'
2023/08/08 17:36:44 INFO mlflow.pyfunc.backend: === Running command 'exec uvicorn --host 127.0.0.1 --port 9020 --workers 1 mlflow.pyfunc.scoring_server.app:app'
INFO: Started server process [6992]
INFO: Waiting for application startup.
INFO: Application startup complete.
INFO: Uvicorn running on http://127.0.0.1:9020
```
The scoring server is ready to receive traffic.
Update the MLflow AI Gateway configuration file (config.yaml) with the new endpoint:
```yaml
endpoints:
- name: embeddings
endpoint_type: llm/v1/embeddings
model:
provider: mlflow-model-serving
name: sentence-transformer
config:
model_server_url: http://127.0.0.1:9020
```
The key component here is the `model_server_url`. For serving an MLflow LLM, this url must match to the service that you are specifying for the
Model Serving server.
> **NOTE:** The MLflow Model Server does not have to be running in order to update the configuration file or to start the MLflow AI Gateway. In order to respond to submitted queries, it is required to be running.
## Creating and logging a fill mask model
To support an additional endpoint for generating a mask fill response from masked input text, we need to log an appropriate model.
For this tutorial example, we'll use a `transformers` `Pipeline` wrapping a `BertForMaskedLM` torch model and will log this pipeline using the MLflow `transformers` flavor.
```python
from transformers import AutoTokenizer, AutoModelForMaskedLM
import mlflow
lm_architecture = "bert-base-cased"
artifact_path = "mask_fill_model"
tokenizer = AutoTokenizer.from_pretrained(lm_architecture)
model = AutoModelForMaskedLM.from_pretrained(lm_architecture)
components = {"model": model, "tokenizer": tokenizer}
with mlflow.start_run():
model_info = mlflow.transformers.log_model(
transformers_model=components,
name=artifact_path,
)
```
## Generate the cli command for starting a local MLflow Model Serving endpoint for this fill mask model
```python
print(f"mlflow models serve -m {model_info.model_uri} -h 127.0.0.1 -p 9010 --no-conda")
```
## Starting the model server for the fill mask model
Using the command printed to stdout from above, open a new terminal (do not close the terminal that is currently running the embeddings model being served!)
and paste the command.
```commandline
mlflow models serve -m file:///Users/me/demos/mlruns/0/bc8bdb7fb90c406eb95603a97742cef8/artifacts/mask_fill_model -h 127.0.0.1 -p 9010 --no-conda
```
## Update the config.yaml to add a new completions endpoint
Ensure that the MLflow serving endpoint starts and is ready for traffic.
```commandline
2023/08/08 17:39:14 INFO mlflow.models.flavor_backend_registry: Selected backend for flavor 'python_function'
2023/08/08 17:39:14 INFO mlflow.pyfunc.backend: === Running command 'exec uvicorn --host 127.0.0.1 --port 9010 --workers 1 mlflow.pyfunc.scoring_server.app:app'
INFO: Started server process [6992]
INFO: Waiting for application startup.
INFO: Application startup complete.
INFO: Uvicorn running on http://127.0.0.1:9010
```
Add the entry to the MLflow AI Gateway configuration file. The final file should match [the config file](config.yaml)
## Create a completions model using MPT-7B-instruct (optional, see notes below)
> **NOTE:** If your system does not have a CUDA-compatible GPU and you have not installed torch with the appropriate CUDA libraries, it is not recommended to attempt to run this portion of the example.
> The inference performance of the MPT-7B-instruct model running on CPU is very slow.
> It is also not recommended to add this model to an MLflow model serving environment that does not have a sufficiently powerful GPU available.
### Download the MPT-7B instruct model and tokenizer to a local directory cache
```python
from huggingface_hub import snapshot_download
snapshot_location = snapshot_download(repo_id="mosaicml/mpt-7b-instruct", local_dir="mpt-7b")
```
### Define the PyFunc model that will be used for the completions endpoint
```python
import transformers
import mlflow
import torch
class MPT(mlflow.pyfunc.PythonModel):
def load_context(self, context):
"""
This method initializes the tokenizer and language model
using the specified model snapshot directory.
"""
# Initialize tokenizer and language model
self.tokenizer = transformers.AutoTokenizer.from_pretrained(
context.artifacts["snapshot"], padding_side="left"
)
config = transformers.AutoConfig.from_pretrained(
context.artifacts["snapshot"], trust_remote_code=True
)
# Comment out this configuration setting if not running on a GPU or if triton is not installed.
# Note that triton dramatically improves the inference speed performance
config.attn_config["attn_impl"] = "triton"
self.model = transformers.AutoModelForCausalLM.from_pretrained(
context.artifacts["snapshot"],
config=config,
torch_dtype=torch.bfloat16,
trust_remote_code=True,
)
# NB: If you do not have a CUDA-capable device or have torch installed with CUDA support
# this setting will not function correctly. Setting device to 'cpu' is valid, but
# the performance will be very slow.
self.model.to(device="cuda")
self.model.eval()
def _build_prompt(self, instruction):
"""
This method generates the prompt for the model.
"""
INSTRUCTION_KEY = "### Instruction:"
RESPONSE_KEY = "### Response:"
INTRO_BLURB = (
"Below is an instruction that describes a task. "
"Write a response that appropriately completes the request."
)
return f"""{INTRO_BLURB}
{INSTRUCTION_KEY}
{instruction}
{RESPONSE_KEY}
"""
def predict(self, context, model_input, params=None):
"""
This method generates prediction for the given input.
"""
prompt = model_input["prompt"][0]
temperature = model_input.get("temperature", [1.0])[0]
max_tokens = model_input.get("max_tokens", [100])[0]
# Build the prompt
prompt = self._build_prompt(prompt)
# Encode the input and generate prediction
# NB: Sending the tokenized inputs to the GPU here explicitly will not work if your system does not have CUDA support.
# If attempting to run this with only CPU support, change 'cuda' to 'cpu'
encoded_input = self.tokenizer.encode(prompt, return_tensors="pt").to("cuda")
output = self.model.generate(
encoded_input,
do_sample=True,
temperature=temperature,
max_new_tokens=max_tokens,
)
# Decode the prediction to text
generated_text = self.tokenizer.decode(output[0], skip_special_tokens=True)
# Removing the prompt from the generated text
prompt_length = len(self.tokenizer.encode(prompt, return_tensors="pt")[0])
generated_response = self.tokenizer.decode(
output[0][prompt_length:], skip_special_tokens=True
)
return {"candidates": [generated_response]}
```
### Specify the model signature, input example, and log the custom model
```python
import pandas as pd
import mlflow
from mlflow.models.signature import ModelSignature
from mlflow.types import DataType, Schema, ColSpec
# Define input and output schema
input_schema = Schema([
ColSpec(DataType.string, "prompt"),
ColSpec(DataType.double, "temperature"),
ColSpec(DataType.long, "max_tokens"),
])
output_schema = Schema([ColSpec(DataType.string, "candidates")])
signature = ModelSignature(inputs=input_schema, outputs=output_schema)
# Define input example
input_example = pd.DataFrame({
"prompt": ["What is machine learning?"],
"temperature": [0.5],
"max_tokens": [100],
})
with mlflow.start_run():
mlflow.pyfunc.log_model(
name="mpt-7b-instruct",
python_model=MPT(),
artifacts={"snapshot": snapshot_location},
pip_requirements=[
"torch",
"transformers",
"accelerate",
"einops",
"sentencepiece",
],
input_example=input_example,
signature=signature,
)
```
## Starting the model server for mpt-7B-instruct (Optional)
Due to the size and complexity of the MPT-7B-instruct model, it is highly advised to only attempt to serve this model in an environment that has:
- A powerful GPU that is capable of holding the model weights in GPU memory
- triton installed
In order to initialize the MLflow Model Server for a large model such as MPT-7B, a slightly modified cli command must be used. Most notably, the timeout duration must be increased from the
default of 60 seconds and it is highly recommended to utilize only a single Gunicorn worker (since each worker will load its own copy of the model, there is a distinct possibility of crashing the server environment with an out of memory fault).
```commandline
mlflow models serve -m file:///Users/me/demos/mlruns/0/92d017e23ca04ffa919a935ed54e9334/artifacts/mpt-7b-instruct -h 127.0.0.1 -p 9030 -t 1200 -w 1 --no-conda
```
## Update the config.yaml to add the MPT-7B-instruct endpoint (Optional)
> **NOTE** If you are adding this endpoint for the example, you will have to manually edit the config.yaml. If the server that is running the MPT-7B-instruct custom PyFunc model's inference does not have GPU support,
> the performance for inference will take a very long time (CPU inference with this model can take tens of minutes for a single query).
```yaml
endpoints:
- name: embeddings
endpoint_type: llm/v1/embeddings
model:
provider: mlflow-model-serving
name: sentence-transformer
config:
model_server_url: http://127.0.0.1:9020
- name: fillmask
endpoint_type: llm/v1/completions
model:
provider: mlflow-model-serving
name: fill-mask
config:
model_server_url: http://127.0.0.1:9010
- name: mpt-instruct
endpoint_type: llm/v1/completions
model:
provider: mlflow-model-serving
name: mpt-7b-instruct
config:
model_server_url: http://127.0.0.1:9030
```
## Start the MLflow AI Gateway
Now that both endpoints (or all 3, if adding in the optional MPT-7B-instruct model endpoint) are defined within the configuration YAML file and the Model Serving servers are ready to receive queries, we can start the MLflow AI Gateway.
```sh
mlflow gateway start --config-path examples/gateway/mlflow_serving/config.yaml --port 7000
```
If adding the mpt-7b-instruct model, start the MLflow AI Gateway by directing the `--config-path` argument to the location of the `config.yaml` file that you've created with the endpoint's addition.
## Query the MLflow AI Gateway
See the [example script](example.py) within this directory to see how to query these two models that are being served.
### Query the mpt-7B-instruct endpoint (Optional)
In order to query the mpt-7b-instruct model, the example shown in the script can be modified by adding an additional query call, as shown below:
```python
# Querying the optional mpt-7b-instruct endpoint
response_mpt = query(
endpoint="mpt-instruct",
data={
"prompt": "What is the purpose of an attention mask in a transformers model?",
"temperature": 0.1,
"max_tokens": 200,
},
)
print(f"Fluent API response for mpt-instruct: {response_mpt}")
```
@@ -0,0 +1,15 @@
endpoints:
- name: fillmask
endpoint_type: llm/v1/completions
model:
provider: mlflow-model-serving
name: mask-fill
config:
model_server_url: http://127.0.0.1:9010
- name: embeddings
endpoint_type: llm/v1/embeddings
model:
provider: mlflow-model-serving
name: sentence-transformer
config:
model_server_url: http://127.0.0.1:9020
+35
View File
@@ -0,0 +1,35 @@
# Prior to running the example code below, view the README.md within this directory
from mlflow.deployments import get_deploy_client
def main():
client = get_deploy_client("http://localhost:7000")
print(f"MLflow model endpoints: {client.list_endpoints()}\n")
print(f"MLflow completions endpoint info: {client.get_endpoint(endpoint='completions')}\n")
# Completions query
response_completions = client.predict(
endpoint="fillmask",
inputs={
"prompt": "I like to [MASK] cars!",
},
)
print(f"MLflow model response for completions: {response_completions}")
# Embeddings query
response_embeddings = client.predict(
endpoint="embeddings",
inputs={
"input"[
"MLflow Deployments sure is useful!",
"Word embeddings are very useful",
]
},
)
print(f"MLflow model response for embeddings: {response_embeddings}")
if __name__ == "__main__":
main()
+13
View File
@@ -0,0 +1,13 @@
## Example endpoint configuration for MosaicML
To see an example of specifying both the completions and the embeddings endpoints for MosaicML, see [the configuration](config.yaml) YAML file.
This configuration file specifies three endpoints: 'completions', 'embeddings', and 'chat', using MosaicML's models 'mpt-7b-instruct', 'instructor-xl', and 'llama2-70b-chat', respectively.
## Setting a MosaicML API Key
This example requires a [MosaicML API key](https://docs.mosaicml.com/en/latest/getting_started.html):
```sh
export MOSAICML_API_KEY=...
```
+24
View File
@@ -0,0 +1,24 @@
endpoints:
- name: completions
endpoint_type: llm/v1/completions
model:
provider: mosaicml
name: mpt-7b-instruct
config:
mosaicml_api_key: $MOSAICML_API_KEY
- name: embeddings
endpoint_type: llm/v1/embeddings
model:
provider: mosaicml
name: instructor-xl
config:
mosaicml_api_key: $MOSAICML_API_KEY
- name: chat
endpoint_type: llm/v1/chat
model:
provider: mosaicml
name: llama2-70b-chat
config:
mosaicml_api_key: $MOSAICML_API_KEY
+48
View File
@@ -0,0 +1,48 @@
from mlflow.deployments import get_deploy_client
def main():
client = get_deploy_client("http://localhost:7000")
print(f"MosaicML endpoints: {client.list_endpoints()}\n")
print(f"MosaicML completions endpoint info: {client.get_endpoint(endpoint='completions')}\n")
# Completions request
response_completions = client.predict(
endpoint="completions",
inputs={
"prompt": "What is the world record for flapjack consumption in a single sitting?",
"temperature": 0.1,
},
)
print(f"MosaicML response for completions: {response_completions}")
# Embeddings request
response_embeddings = client.predict(
endpoint="embeddings",
inputs={"input": ["Do you carry the Storm Trooper costume in size 2T?"]},
)
print(f"MosaicML response for embeddings: {response_embeddings}")
# Chat example
response_chat = client.predict(
endpoint="chat",
inputs={
"messages": [
{
"role": "system",
"content": "You are a talented European rapper with a background in US history",
},
{
"role": "user",
"content": "Please recite the preamble to the US Constitution as if it were "
"written today by a rapper from Reykjavík",
},
]
},
)
print(f"MosaicML response for chat: {response_chat}")
if __name__ == "__main__":
main()
+15
View File
@@ -0,0 +1,15 @@
## Example endpoint configuration for OpenAI
To view an example of OpenAI endpoint configurations, see [the configuration example](config.yaml) YAML file for OpenAI.
This configuration shows all 3 supported endpoint types: chat, completions, and embeddings.
## Setting the OpenAI API Key
An OpenAI API key is required for the configuration. If you haven't already, obtain an [OpenAI API key](https://platform.openai.com/account/api-keys).
With the key, export it to your environment variables. Replace the '...' with your actual API key:
```sh
export OPENAI_API_KEY=...
```
+27
View File
@@ -0,0 +1,27 @@
endpoints:
- name: chat
endpoint_type: llm/v1/chat
model:
provider: openai
name: gpt-4o-mini
config:
openai_api_key: $OPENAI_API_KEY
limit:
renewal_period: minute
calls: 10
- name: completions
endpoint_type: llm/v1/completions
model:
provider: openai
name: gpt-4o-mini
config:
openai_api_key: $OPENAI_API_KEY
- name: embeddings
endpoint_type: llm/v1/embeddings
model:
provider: openai
name: text-embedding-ada-002
config:
openai_api_key: $OPENAI_API_KEY
+47
View File
@@ -0,0 +1,47 @@
from mlflow.deployments import get_deploy_client
def main():
client = get_deploy_client("http://localhost:7000")
print(f"OpenAI endpoints: {client.list_endpoints()}\n")
print(f"OpenAI endpoint info: {client.get_endpoint(endpoint='completions')}\n")
# Completions example
response_completions = client.predict(
endpoint="completions",
inputs={
"prompt": "How many patties could be stacked on a cheeseburger before issues arise?",
"max_tokens": 200,
"temperature": 0.25,
},
)
print(f"OpenAI completions response: {response_completions}")
# Chat example
response_chat = client.predict(
endpoint="chat",
inputs={
"messages": [
{
"role": "user",
"content": "Please recite the preamble to the US Constitution as if it were "
"written today by a rapper from Reykjavík",
}
]
},
)
print(f"OpenAI completions response: {response_chat}")
# Embeddings example
response_embeddings = client.predict(
endpoint="embeddings",
inputs={
"input": "When you say 'enriched', what exactly are you enriching the cereal with?"
},
)
print(f"OpenAI response for embeddings: {response_embeddings}")
if __name__ == "__main__":
main()
+13
View File
@@ -0,0 +1,13 @@
## Example endpoint configuration for PaLM
To see an example of specifying both the completions and the embeddings endpoints for PaLM, see [the configuration](config.yaml) YAML file.
This configuration file specifies three endpoints: 'completions', 'embeddings', and 'chat', using PaLM's models 'text-bison-001', 'embedding-gecko-001', and 'chat-bison-001', respectively.
## Setting a PaLM API Key
This example requires a [PaLM API key](https://developers.generativeai.google/tutorials/setup):
```sh
export PALM_API_KEY=...
```
+24
View File
@@ -0,0 +1,24 @@
endpoints:
- name: completions
endpoint_type: llm/v1/completions
model:
provider: palm
name: text-bison-001
config:
palm_api_key: $PALM_API_KEY
- name: embeddings
endpoint_type: llm/v1/embeddings
model:
provider: palm
name: embedding-gecko-001
config:
palm_api_key: $PALM_API_KEY
- name: chat
endpoint_type: llm/v1/chat
model:
provider: palm
name: chat-bison-001
config:
palm_api_key: $PALM_API_KEY
+48
View File
@@ -0,0 +1,48 @@
from mlflow.deployments import get_deploy_client
def main():
client = get_deploy_client("http://localhost:7000")
print(f"PaLM endpoints: {client.list_endpoints()}\n")
print(f"PaLM completions endpoint info: {client.get_endpoint(endpoint='completions')}\n")
# Completions request
response_completions = client.predict(
endpoint="completions",
inputs={
"prompt": "What is the world record for flapjack consumption in a single sitting?",
"temperature": 0.1,
},
)
print(f"PaLM response for completions: {response_completions}")
# Embeddings request
response_embeddings = client.predict(
endpoint="embeddings",
inputs={"input": ["Do you carry the Storm Trooper costume in size 2T?"]},
)
print(f"PaLM response for embeddings: {response_embeddings}")
# Chat example
response_chat = client.predict(
endpoint="chat",
inputs={
"messages": [
{
"role": "system",
"content": "You are a talented European rapper with a background in US history",
},
{
"role": "user",
"content": "Please recite the preamble to the US Constitution as if it were "
"written today by a rapper from Reykjavík",
},
]
},
)
print(f"PaLM response for chat: {response_chat}")
if __name__ == "__main__":
main()
+28
View File
@@ -0,0 +1,28 @@
## Example endpoint configuration for plugin provider
To see an example of specifying the chat endpoint for a plugin provider,
see [the configuration](config.yaml) YAML file.
We implement our plugin provider package `my_llm` under `./my-llm` folder. It implements the chat method.
This configuration file specifies one endpoint: 'chat', using the model 'my-model-0.1.2'.
## Setting up the server
First, install the provider package `my_llm`:
```sh
pip install -e ./my-llm
```
Then, start the server:
```sh
MY_LLM_API_KEY=some-api-key mlflow gateway start --config-path config.yaml --port 7000
```
To clean up the installed package after the example, run
```sh
pip uninstall my_llm
```
+8
View File
@@ -0,0 +1,8 @@
endpoints:
- name: chat
endpoint_type: llm/v1/chat
model:
provider: my_llm
name: my-model-0.1.2
config:
my_llm_api_key: $MY_LLM_API_KEY
+26
View File
@@ -0,0 +1,26 @@
from mlflow.deployments import get_deploy_client
def main():
client = get_deploy_client("http://127.0.0.1:7000")
print(f"Plugin endpoints: {client.list_endpoints()}\n")
print(f"Plugin chat endpoint info: {client.get_endpoint(endpoint='chat')}\n")
# Chat request
response_chat = client.predict(
endpoint="chat",
inputs={
"messages": [
{
"role": "user",
"content": "Tell me a joke",
}
]
},
)
print(f"Plugin response for chat: {response_chat}")
if __name__ == "__main__":
main()
@@ -0,0 +1,20 @@
import os
from pydantic import field_validator
from mlflow.gateway.base_models import ConfigModel
class MyLLMConfig(ConfigModel):
my_llm_api_key: str
@field_validator("my_llm_api_key", mode="before")
def validate_my_llm_api_key(cls, value):
if value.startswith("$"):
# This resolves the API key from an environment variable
env_var_name = value[1:]
if env_var := os.environ.get(env_var_name):
return env_var
else:
raise ValueError(f"Environment variable {env_var_name!r} is not set")
return value
@@ -0,0 +1,37 @@
import time
from mlflow.gateway.config import EndpointConfig
from mlflow.gateway.providers import BaseProvider
from mlflow.gateway.schemas import chat
from my_llm.config import MyLLMConfig
class MyLLMProvider(BaseProvider):
NAME = "MyLLM"
CONFIG_TYPE = MyLLMConfig
def __init__(self, config: EndpointConfig) -> None:
super().__init__(config)
if config.model.config is None or not isinstance(config.model.config, MyLLMConfig):
raise TypeError(f"Unexpected config type {config.model.config}")
self.my_llm_config: MyLLMConfig = config.model.config
async def chat(self, payload: chat.RequestPayload) -> chat.ResponsePayload:
return chat.ResponsePayload(
id="id-123",
created=int(time.time()),
model=self.config.model.name,
choices=[
chat.Choice(
index=0,
message=chat.ResponseMessage(
role="assistant", content="This is a response from MyLLMProvider"
),
)
],
usage=chat.ChatUsage(
prompt_tokens=10,
completion_tokens=18,
total_tokens=28,
),
)
@@ -0,0 +1,10 @@
[project]
name = "my_llm"
version = "1.0"
[project.entry-points."mlflow.gateway.providers"]
my_llm = "my_llm.providers:MyLLMProvider"
[tool.setuptools.packages.find]
include = ["my_llm*"]
namespaces = false
+13
View File
@@ -0,0 +1,13 @@
## Example endpoint configuration for TogetherAI
To see an example of specifying both the completions and the embeddings endpoints for TogetherAI, see [the configuration](config.yaml) YAML file.
This configuration file specifies two endpoints: 'completions' and 'embeddings', both using TogetherAI's provided models 'mistralai/Mixtral-8x7B-v0.1' and 'togethercomputer/m2-bert-80M-8k-retrieval', respectively.
## Setting a Mistral API Key
This example requires a [TogetherAI API key](https://docs.together.ai/docs/):
```sh
export TOGETHERAI_API_KEY=...
```
+24
View File
@@ -0,0 +1,24 @@
endpoints:
- name: completions
endpoint_type: llm/v1/completions
model:
provider: togetherai
name: mistralai/Mixtral-8x7B-v0.1
config:
togetherai_api_key: $TOGETHERAI_API_KEY
- name: chat
endpoint_type: llm/v1/chat
model:
provider: togetherai
name: mistralai/Mixtral-8x7B-Instruct-v0.1
config:
togetherai_api_key: $TOGETHERAI_API_KEY
- name: embeddings
endpoint_type: llm/v1/embeddings
model:
provider: togetherai
name: togethercomputer/m2-bert-80M-8k-retrieval
config:
togetherai_api_key: $TOGETHERAI_API_KEY
+43
View File
@@ -0,0 +1,43 @@
from mlflow.deployments import get_deploy_client
def main():
client = get_deploy_client("http://localhost:7000")
print(f"Togetherai endpoints: {client.list_endpoints()}\n")
print(f"Togetherai completions endpoint info: {client.get_endpoint(endpoint='completions')}\n")
print(f"Togetherai chat endpoint info: {client.get_endpoint(endpoint='chat')}\n")
print(f"Togetherai embeddings endpoint info: {client.get_endpoint(endpoint='embeddings')}\n")
response_completions = client.predict(
endpoint="completions",
inputs={
"prompt": "Who is the protagonist in Witcher 3 Wild Hunt?",
"max_tokens": 200,
"temperature": 0.1,
},
)
print(f"Togetherai response for completions: {response_completions}")
response_embeddings = client.predict(
endpoint="embeddings",
inputs={
"input": ["Who is Wes Montgomery?"],
},
)
print(f"Togetherai response for embeddings: {response_embeddings}")
response_chat = client.predict(
endpoint="chat",
inputs={
"messages": [{"role": "user", "content": "Get out of the sunlight's way Alexander!"}],
},
)
print(f"Togetherai response for chat: {response_chat}")
if __name__ == "__main__":
main()
+59
View File
@@ -0,0 +1,59 @@
# Unity Catalog Integration
This example demonstrates how to use the Unity Catalog (UC) integration with MLflow AI Gateway.
## Pre-requisites
1. Install the required packages:
```bash
pip install mlflow openai databricks-sdk
```
2. Create the UC function used in `run.py` by running the following command on Databricks notebook:
```
%sql
CREATE OR REPLACE FUNCTION
my.uc_func.add (
x INTEGER COMMENT 'The first number to add.',
y INTEGER COMMENT 'The second number to add.'
)
RETURNS INTEGER
LANGUAGE SQL
RETURN x + y
```
To define your own function, see https://docs.databricks.com/en/sql/language-manual/sql-ref-syntax-ddl-create-sql-function.html#create-function-sql-and-python.
3. Create a SQL warehouse in Databricks by following the instructions at https://docs.databricks.com/en/compute/sql-warehouse/create.html.
## Running the example script
First, run the deployments server:
```bash
# Required to authenticate with Databricks. See https://docs.databricks.com/en/dev-tools/auth/index.html#supported-authentication-types-by-databricks-tool-or-sdk for other authentication methods.
export DATABRICKS_HOST="..." # e.g. https://my.databricks.com
export DATABRICKS_TOKEN="..."
# Required to execute UC functions. See https://docs.databricks.com/en/integrations/compute-details.html#get-connection-details-for-a-databricks-compute-resource for how to get the http path of your warehouse.
# The last part of the http path is the warehouse ID.
#
# /sql/1.0/warehouses/1234567890123456
# ^^^^^^^^^^^^^^^^
export DATABRICKS_WAREHOUSE_ID="..."
# Enable Unity Catalog integration
export MLFLOW_ENABLE_UC_FUNCTIONS=true
mlflow gateway start --config-path examples/gateway/openai/config.yaml --port 7000
```
Once the server starts running, run the example script:
```bash
# Replace `my.uc_func.add` if your UC function has a different name
python examples/gateway/uc_functions/run.py --uc-function-name my.uc_func.add
```
+113
View File
@@ -0,0 +1,113 @@
import argparse
import json
import openai
def parse_args():
parser = argparse.ArgumentParser()
parser.add_argument(
"--uc-function-name",
type=str,
required=True,
help="Name of the UC function to use",
)
return parser.parse_args()
def main():
args = parse_args()
client = openai.OpenAI(base_url="http://localhost:7000/v1")
print("----- UC function -----")
uc_function = {
"type": "uc_function",
"uc_function": {
"name": args.uc_function_name,
},
}
resp = client.chat.completions.create(
model="chat",
messages=[
{
"role": "user",
"content": "What is the result of 1 + 2?",
}
],
tools=[uc_function],
)
print(resp.choices[0].message.content)
print("----- UC function + User-defined function -----")
user_defined_function = {
"type": "function",
"function": {
"description": "Multiply numbers",
"name": "multiply",
"parameters": {
"type": "object",
"properties": {
"x": {
"type": "integer",
"description": "First number",
},
"y": {
"type": "integer",
"description": "Second number",
},
},
"required": ["x", "y"],
},
},
}
def multiply(x: int, y: int) -> int:
return x * y
msg = {
"role": "user",
"content": (
"What is the result of 1 + 2? What is the result of 3 + 4? What is the result of 5 * 6?"
),
}
resp = client.chat.completions.create(
model="chat",
messages=[msg],
tools=[
user_defined_function,
uc_function,
],
)
print(resp.choices[0].message.content)
print(resp.choices[0].message.tool_calls)
multiply_call = resp.choices[0].message.tool_calls[0].function
assert multiply_call.name == "multiply"
resp = client.chat.completions.create(
model="chat",
messages=[
msg,
{
"role": "assistant",
"content": resp.choices[0].message.content,
},
{
"role": "assistant",
"content": "",
"tool_calls": resp.choices[0].message.tool_calls,
},
{
"role": "tool",
"tool_call_id": resp.choices[0].message.tool_calls[0].id,
"content": str(multiply(**json.loads(multiply_call.arguments))),
},
],
)
print(resp.choices[0].message.content)
if __name__ == "__main__":
main()
+39
View File
@@ -0,0 +1,39 @@
"""
This is an example for leveraging MLflow's auto tracing capabilities for Gemini.
For more information about MLflow Tracing, see: https://mlflow.org/docs/latest/llms/tracing/index.html
"""
import os
import mlflow
# Turn on auto tracing for Gemini by calling mlflow.gemini.autolog()
mlflow.gemini.autolog()
# Import the SDK and configure your API key.
from google import genai
client = genai.Client(api_key=os.environ["GEMINI_API_KEY"])
# Use the generate_content method to generate responses to your prompts.
response = client.models.generate_content(
model="gemini-1.5-flash", contents="The opposite of hot is"
)
print(response.text)
# Also leverage the chat feature to conduct multi-turn interactions
chat = client.chats.create(model="gemini-1.5-flash")
response = chat.send_message("In one sentence, explain how a computer works to a young child.")
print(response.text)
response = chat.send_message("Okay, how about a more detailed explanation to a high schooler?")
print(response.text)
# Count tokens for your statement
response = client.models.count_tokens("The quick brown fox jumps over the lazy dog.")
print(response.total_tokens)
# Generate text embeddings for your content
text = "Hello world"
result = client.models.embed_content(model="text-embedding-004", contents=text)
print(result["embedding"])
+27
View File
@@ -0,0 +1,27 @@
"""
This is an example for leveraging MLflow's auto tracing capabilities for Groq.
For more information about MLflow Tracing, see: https://mlflow.org/docs/latest/llms/tracing/index.html
"""
import groq
import mlflow
# Turn on auto tracing for Groq by calling mlflow.groq.autolog()
mlflow.groq.autolog()
client = groq.Groq()
# Use the create method to create new message
message = client.chat.completions.create(
model="llama3-8b-8192",
messages=[
{
"role": "user",
"content": "Explain the importance of low latency LLMs.",
}
],
)
print(message.choices[0].message.content)

Some files were not shown because too many files have changed in this diff Show More