chore: import upstream snapshot with attribution
This commit is contained in:
@@ -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.
|
||||
@@ -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?")
|
||||
@@ -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,
|
||||
)
|
||||
@@ -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)
|
||||
@@ -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'
|
||||
```
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -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
|
||||
```
|
||||
@@ -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")
|
||||
@@ -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))
|
||||
@@ -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
|
||||
}
|
||||
@@ -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()
|
||||
@@ -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 .
|
||||
@@ -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}"
|
||||
@@ -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.
|
||||
@@ -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
|
||||
@@ -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
@@ -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
@@ -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
|
||||
}
|
||||
@@ -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}"
|
||||
|
||||
@@ -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
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -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
|
||||
```
|
||||
@@ -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=...
|
||||
```
|
||||
@@ -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
|
||||
@@ -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()
|
||||
@@ -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=...
|
||||
```
|
||||
@@ -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
|
||||
@@ -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()
|
||||
@@ -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.
|
||||
@@ -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"
|
||||
@@ -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.
|
||||
@@ -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
|
||||
@@ -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()
|
||||
@@ -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=...
|
||||
```
|
||||
@@ -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
|
||||
@@ -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()
|
||||
@@ -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=...
|
||||
```
|
||||
@@ -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
|
||||
@@ -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()
|
||||
@@ -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.
|
||||
@@ -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
|
||||
@@ -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()
|
||||
@@ -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=...
|
||||
```
|
||||
@@ -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
|
||||
@@ -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()
|
||||
@@ -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
|
||||
@@ -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()
|
||||
@@ -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=...
|
||||
```
|
||||
@@ -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
|
||||
@@ -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()
|
||||
@@ -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=...
|
||||
```
|
||||
@@ -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
|
||||
@@ -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()
|
||||
@@ -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=...
|
||||
```
|
||||
@@ -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
|
||||
@@ -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()
|
||||
@@ -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
|
||||
```
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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=...
|
||||
```
|
||||
@@ -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
|
||||
@@ -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()
|
||||
@@ -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
|
||||
```
|
||||
@@ -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()
|
||||
@@ -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"])
|
||||
@@ -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
Reference in New Issue
Block a user