wehub-resource-sync
662 lines
21 KiB
Python
662 lines
21 KiB
Python
import json
|
|
import pathlib
|
|
import pickle
|
|
import uuid
|
|
from dataclasses import asdict
|
|
|
|
import pytest
|
|
|
|
import mlflow
|
|
from mlflow.exceptions import MlflowException
|
|
from mlflow.models.model import Model
|
|
from mlflow.models.signature import ModelSignature
|
|
from mlflow.models.utils import load_serving_example
|
|
from mlflow.pyfunc.loaders.chat_model import _ChatModelPyfuncWrapper
|
|
from mlflow.tracing.constant import TraceTagKey
|
|
from mlflow.tracking.artifact_utils import _download_artifact_from_uri
|
|
from mlflow.types.llm import (
|
|
CHAT_MODEL_INPUT_SCHEMA,
|
|
CHAT_MODEL_OUTPUT_SCHEMA,
|
|
ChatChoice,
|
|
ChatChoiceDelta,
|
|
ChatChunkChoice,
|
|
ChatCompletionChunk,
|
|
ChatCompletionResponse,
|
|
ChatMessage,
|
|
ChatParams,
|
|
FunctionToolCallArguments,
|
|
FunctionToolDefinition,
|
|
ToolParamsSchema,
|
|
)
|
|
from mlflow.types.schema import ColSpec, DataType, Schema
|
|
|
|
from tests.helper_functions import (
|
|
expect_status_code,
|
|
pyfunc_serve_and_score_model,
|
|
)
|
|
from tests.tracing.helper import get_traces
|
|
|
|
# `None`s (`max_tokens` and `stop`) are excluded
|
|
DEFAULT_PARAMS = {
|
|
"temperature": 1.0,
|
|
"n": 1,
|
|
"stream": False,
|
|
}
|
|
|
|
|
|
def get_mock_streaming_response(message, is_last_chunk=False):
|
|
if is_last_chunk:
|
|
return {
|
|
"id": "123",
|
|
"model": "MyChatModel",
|
|
"choices": [
|
|
{
|
|
"index": 0,
|
|
"delta": {
|
|
"role": None,
|
|
"content": None,
|
|
},
|
|
"finish_reason": "stop",
|
|
}
|
|
],
|
|
"usage": {
|
|
"prompt_tokens": 10,
|
|
"completion_tokens": 10,
|
|
"total_tokens": 20,
|
|
},
|
|
}
|
|
else:
|
|
return {
|
|
"id": "123",
|
|
"model": "MyChatModel",
|
|
"choices": [
|
|
{
|
|
"index": 0,
|
|
"delta": {
|
|
"role": "assistant",
|
|
"content": message,
|
|
},
|
|
"finish_reason": "stop",
|
|
}
|
|
],
|
|
"usage": {
|
|
"prompt_tokens": 10,
|
|
"completion_tokens": 10,
|
|
"total_tokens": 20,
|
|
},
|
|
}
|
|
|
|
|
|
def get_mock_response(messages, params):
|
|
return {
|
|
"id": "123",
|
|
"model": "MyChatModel",
|
|
"choices": [
|
|
{
|
|
"index": 0,
|
|
"message": {
|
|
"role": "assistant",
|
|
"content": json.dumps([m.to_dict() for m in messages]),
|
|
},
|
|
"finish_reason": "stop",
|
|
},
|
|
{
|
|
"index": 1,
|
|
"message": {
|
|
"role": "user",
|
|
"content": json.dumps(params.to_dict()),
|
|
},
|
|
"finish_reason": "stop",
|
|
},
|
|
],
|
|
"usage": {
|
|
"prompt_tokens": 10,
|
|
"completion_tokens": 10,
|
|
"total_tokens": 20,
|
|
},
|
|
}
|
|
|
|
|
|
class SimpleChatModel(mlflow.pyfunc.ChatModel):
|
|
def predict(
|
|
self, context, messages: list[ChatMessage], params: ChatParams
|
|
) -> ChatCompletionResponse:
|
|
mock_response = get_mock_response(messages, params)
|
|
return ChatCompletionResponse.from_dict(mock_response)
|
|
|
|
def predict_stream(self, context, messages: list[ChatMessage], params: ChatParams):
|
|
num_chunks = 10
|
|
for i in range(num_chunks):
|
|
mock_response = get_mock_streaming_response(
|
|
f"message {i}", is_last_chunk=(i == num_chunks - 1)
|
|
)
|
|
yield ChatCompletionChunk.from_dict(mock_response)
|
|
|
|
|
|
class ChatModelWithContext(mlflow.pyfunc.ChatModel):
|
|
def load_context(self, context):
|
|
predict_path = pathlib.Path(context.artifacts["predict_fn"])
|
|
self.predict_fn = pickle.loads(predict_path.read_bytes())
|
|
|
|
def predict(
|
|
self, context, messages: list[ChatMessage], params: ChatParams
|
|
) -> ChatCompletionResponse:
|
|
message = ChatMessage(role="assistant", content=self.predict_fn())
|
|
return ChatCompletionResponse.from_dict(get_mock_response([message], params))
|
|
|
|
|
|
class ChatModelWithTrace(mlflow.pyfunc.ChatModel):
|
|
@mlflow.trace
|
|
def predict(
|
|
self, context, messages: list[ChatMessage], params: ChatParams
|
|
) -> ChatCompletionResponse:
|
|
mock_response = get_mock_response(messages, params)
|
|
return ChatCompletionResponse.from_dict(mock_response)
|
|
|
|
|
|
class ChatModelWithMetadata(mlflow.pyfunc.ChatModel):
|
|
def predict(
|
|
self, context, messages: list[ChatMessage], params: ChatParams
|
|
) -> ChatCompletionResponse:
|
|
mock_response = get_mock_response(messages, params)
|
|
return ChatCompletionResponse(
|
|
**mock_response,
|
|
custom_outputs=params.custom_inputs,
|
|
)
|
|
|
|
|
|
class ChatModelWithToolCalling(mlflow.pyfunc.ChatModel):
|
|
def predict(
|
|
self, context, messages: list[ChatMessage], params: ChatParams
|
|
) -> ChatCompletionResponse:
|
|
tools = params.tools
|
|
|
|
# call the first tool with some value for all the required params
|
|
tool_name = tools[0].function.name
|
|
tool_params = tools[0].function.parameters
|
|
arguments = {}
|
|
for param in tool_params.required:
|
|
param_type = tool_params.properties[param].type
|
|
if param_type == "string":
|
|
arguments[param] = "some_value"
|
|
elif param_type == "number":
|
|
arguments[param] = 123
|
|
elif param_type == "boolean":
|
|
arguments[param] = True
|
|
else:
|
|
# keep the test example simple
|
|
raise ValueError(f"Unsupported param type: {param_type}")
|
|
|
|
tool_call = FunctionToolCallArguments(
|
|
name=tool_name,
|
|
arguments=json.dumps(arguments),
|
|
).to_tool_call(id=uuid.uuid4().hex)
|
|
|
|
tool_message = ChatMessage(
|
|
role="assistant",
|
|
tool_calls=[tool_call],
|
|
)
|
|
|
|
return ChatCompletionResponse(choices=[ChatChoice(index=0, message=tool_message)])
|
|
|
|
|
|
def test_chat_model_save_load(tmp_path):
|
|
model = SimpleChatModel()
|
|
mlflow.pyfunc.save_model(python_model=model, path=tmp_path)
|
|
|
|
loaded_model = mlflow.pyfunc.load_model(tmp_path)
|
|
assert isinstance(loaded_model._model_impl, _ChatModelPyfuncWrapper)
|
|
input_schema = loaded_model.metadata.get_input_schema()
|
|
output_schema = loaded_model.metadata.get_output_schema()
|
|
assert input_schema == CHAT_MODEL_INPUT_SCHEMA
|
|
assert output_schema == CHAT_MODEL_OUTPUT_SCHEMA
|
|
|
|
|
|
def test_chat_model_with_trace(tmp_path):
|
|
model = ChatModelWithTrace()
|
|
mlflow.pyfunc.save_model(python_model=model, path=tmp_path)
|
|
|
|
# predict() call during saving chat model should not generate a trace
|
|
assert len(get_traces()) == 0
|
|
|
|
loaded_model = mlflow.pyfunc.load_model(tmp_path)
|
|
messages = [
|
|
{"role": "system", "content": "You are a helpful assistant"},
|
|
{"role": "user", "content": "Hello!"},
|
|
]
|
|
loaded_model.predict({"messages": messages})
|
|
|
|
traces = get_traces()
|
|
assert len(traces) == 1
|
|
assert traces[0].info.tags[TraceTagKey.TRACE_NAME] == "predict"
|
|
request = json.loads(traces[0].data.request)
|
|
assert request["messages"] == [asdict(ChatMessage.from_dict(msg)) for msg in messages]
|
|
|
|
|
|
def test_chat_model_save_throws_with_signature(tmp_path):
|
|
model = SimpleChatModel()
|
|
|
|
with pytest.raises(MlflowException, match="Please remove the `signature` parameter"):
|
|
mlflow.pyfunc.save_model(
|
|
python_model=model,
|
|
path=tmp_path,
|
|
signature=ModelSignature(
|
|
Schema([ColSpec(name="test", type=DataType.string)]),
|
|
Schema([ColSpec(name="test", type=DataType.string)]),
|
|
),
|
|
)
|
|
|
|
|
|
def mock_predict():
|
|
return "hello"
|
|
|
|
|
|
def test_chat_model_with_context_saves_successfully(tmp_path):
|
|
model_path = tmp_path / "model"
|
|
predict_path = tmp_path / "predict.pkl"
|
|
predict_path.write_bytes(pickle.dumps(mock_predict))
|
|
|
|
model = ChatModelWithContext()
|
|
mlflow.pyfunc.save_model(
|
|
python_model=model,
|
|
path=model_path,
|
|
artifacts={"predict_fn": str(predict_path)},
|
|
)
|
|
|
|
loaded_model = mlflow.pyfunc.load_model(model_path)
|
|
messages = [{"role": "user", "content": "test"}]
|
|
|
|
response = loaded_model.predict({"messages": messages})
|
|
expected_response = json.dumps([{"role": "assistant", "content": "hello"}])
|
|
assert response["choices"][0]["message"]["content"] == expected_response
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"ret",
|
|
[
|
|
"not a ChatCompletionResponse",
|
|
{"dict": "with", "bad": "keys"},
|
|
{
|
|
"id": "1",
|
|
"created": 1,
|
|
"model": "m",
|
|
"choices": [{"bad": "choice"}],
|
|
"usage": {
|
|
"prompt_tokens": 10,
|
|
"completion_tokens": 10,
|
|
"total_tokens": 20,
|
|
},
|
|
},
|
|
],
|
|
)
|
|
def test_save_throws_on_invalid_output(tmp_path, ret):
|
|
class BadChatModel(mlflow.pyfunc.ChatModel):
|
|
def predict(self, context, messages, params) -> ChatCompletionResponse:
|
|
return ret
|
|
|
|
model = BadChatModel()
|
|
with pytest.raises(
|
|
MlflowException,
|
|
match=(
|
|
"Failed to save ChatModel. Please ensure that the model's "
|
|
r"predict\(\) method returns a ChatCompletionResponse object"
|
|
),
|
|
):
|
|
mlflow.pyfunc.save_model(python_model=model, path=tmp_path)
|
|
|
|
|
|
# test that we can predict with the model
|
|
def test_chat_model_predict(tmp_path):
|
|
model = SimpleChatModel()
|
|
mlflow.pyfunc.save_model(python_model=model, path=tmp_path)
|
|
|
|
loaded_model = mlflow.pyfunc.load_model(tmp_path)
|
|
messages = [
|
|
{"role": "system", "content": "You are a helpful assistant"},
|
|
{"role": "user", "content": "Hello!"},
|
|
]
|
|
|
|
response = loaded_model.predict({"messages": messages})
|
|
assert response["choices"][0]["message"]["content"] == json.dumps(messages)
|
|
assert json.loads(response["choices"][1]["message"]["content"]) == DEFAULT_PARAMS
|
|
|
|
# override all params
|
|
params_override = {
|
|
"temperature": 0.5,
|
|
"max_tokens": 10,
|
|
"stop": ["\n"],
|
|
"n": 2,
|
|
"stream": True,
|
|
"top_p": 0.1,
|
|
"top_k": 20,
|
|
"frequency_penalty": 0.5,
|
|
"presence_penalty": -0.5,
|
|
}
|
|
response = loaded_model.predict({"messages": messages, **params_override})
|
|
assert response["choices"][0]["message"]["content"] == json.dumps(messages)
|
|
assert json.loads(response["choices"][1]["message"]["content"]) == params_override
|
|
|
|
# override a subset of params
|
|
params_subset = {
|
|
"max_tokens": 100,
|
|
}
|
|
response = loaded_model.predict({"messages": messages, **params_subset})
|
|
assert response["choices"][0]["message"]["content"] == json.dumps(messages)
|
|
assert json.loads(response["choices"][1]["message"]["content"]) == {
|
|
**DEFAULT_PARAMS,
|
|
**params_subset,
|
|
}
|
|
|
|
|
|
def test_chat_model_works_in_serving():
|
|
model = SimpleChatModel()
|
|
messages = [
|
|
{"role": "system", "content": "You are a helpful assistant"},
|
|
{"role": "user", "content": "Hello!"},
|
|
]
|
|
params_subset = {
|
|
"max_tokens": 100,
|
|
}
|
|
with mlflow.start_run():
|
|
model_info = mlflow.pyfunc.log_model(
|
|
name="model",
|
|
python_model=model,
|
|
input_example=(messages, params_subset),
|
|
)
|
|
|
|
inference_payload = load_serving_example(model_info.model_uri)
|
|
response = pyfunc_serve_and_score_model(
|
|
model_uri=model_info.model_uri,
|
|
data=inference_payload,
|
|
content_type="application/json",
|
|
extra_args=["--env-manager", "local"],
|
|
)
|
|
|
|
expect_status_code(response, 200)
|
|
choices = json.loads(response.content)["choices"]
|
|
assert choices[0]["message"]["content"] == json.dumps(messages)
|
|
assert json.loads(choices[1]["message"]["content"]) == {
|
|
**DEFAULT_PARAMS,
|
|
**params_subset,
|
|
}
|
|
|
|
|
|
def test_chat_model_works_with_infer_signature_input_example(tmp_path):
|
|
model = SimpleChatModel()
|
|
params_subset = {
|
|
"max_tokens": 100,
|
|
}
|
|
input_example = {
|
|
"messages": [
|
|
{
|
|
"role": "user",
|
|
"content": "What is Retrieval-augmented Generation?",
|
|
}
|
|
],
|
|
**params_subset,
|
|
}
|
|
with mlflow.start_run():
|
|
model_info = mlflow.pyfunc.log_model(
|
|
name="model", python_model=model, input_example=input_example
|
|
)
|
|
assert model_info.signature.inputs == CHAT_MODEL_INPUT_SCHEMA
|
|
assert model_info.signature.outputs == CHAT_MODEL_OUTPUT_SCHEMA
|
|
mlflow_model = Model.load(model_info.model_uri)
|
|
local_path = _download_artifact_from_uri(model_info.model_uri)
|
|
assert mlflow_model.load_input_example(local_path) == {
|
|
"messages": input_example["messages"],
|
|
**params_subset,
|
|
}
|
|
|
|
inference_payload = load_serving_example(model_info.model_uri)
|
|
response = pyfunc_serve_and_score_model(
|
|
model_uri=model_info.model_uri,
|
|
data=inference_payload,
|
|
content_type="application/json",
|
|
extra_args=["--env-manager", "local"],
|
|
)
|
|
|
|
expect_status_code(response, 200)
|
|
choices = json.loads(response.content)["choices"]
|
|
assert choices[0]["message"]["content"] == json.dumps(input_example["messages"])
|
|
assert json.loads(choices[1]["message"]["content"]) == {
|
|
**DEFAULT_PARAMS,
|
|
**params_subset,
|
|
}
|
|
|
|
|
|
def test_chat_model_logs_default_metadata_task(tmp_path):
|
|
model = SimpleChatModel()
|
|
params_subset = {
|
|
"max_tokens": 100,
|
|
}
|
|
input_example = {
|
|
"messages": [
|
|
{
|
|
"role": "user",
|
|
"content": "What is Retrieval-augmented Generation?",
|
|
}
|
|
],
|
|
**params_subset,
|
|
}
|
|
with mlflow.start_run():
|
|
model_info = mlflow.pyfunc.log_model(
|
|
name="model", python_model=model, input_example=input_example
|
|
)
|
|
assert model_info.signature.inputs == CHAT_MODEL_INPUT_SCHEMA
|
|
assert model_info.signature.outputs == CHAT_MODEL_OUTPUT_SCHEMA
|
|
assert model_info.metadata["task"] == "agent/v1/chat"
|
|
|
|
with mlflow.start_run():
|
|
model_info_with_override = mlflow.pyfunc.log_model(
|
|
name="model", python_model=model, input_example=input_example, metadata={"task": None}
|
|
)
|
|
assert model_info_with_override.metadata["task"] is None
|
|
|
|
|
|
def test_chat_model_works_with_chat_message_input_example(tmp_path):
|
|
model = SimpleChatModel()
|
|
input_example = [
|
|
ChatMessage(role="user", content="What is Retrieval-augmented Generation?", name="chat")
|
|
]
|
|
with mlflow.start_run():
|
|
model_info = mlflow.pyfunc.log_model(
|
|
name="model", python_model=model, input_example=input_example
|
|
)
|
|
assert model_info.signature.inputs == CHAT_MODEL_INPUT_SCHEMA
|
|
assert model_info.signature.outputs == CHAT_MODEL_OUTPUT_SCHEMA
|
|
mlflow_model = Model.load(model_info.model_uri)
|
|
local_path = _download_artifact_from_uri(model_info.model_uri)
|
|
assert mlflow_model.load_input_example(local_path) == {
|
|
"messages": [message.to_dict() for message in input_example],
|
|
}
|
|
|
|
inference_payload = load_serving_example(model_info.model_uri)
|
|
response = pyfunc_serve_and_score_model(
|
|
model_uri=model_info.model_uri,
|
|
data=inference_payload,
|
|
content_type="application/json",
|
|
extra_args=["--env-manager", "local"],
|
|
)
|
|
|
|
expect_status_code(response, 200)
|
|
choices = json.loads(response.content)["choices"]
|
|
assert choices[0]["message"]["content"] == json.dumps(json.loads(inference_payload)["messages"])
|
|
|
|
|
|
def test_chat_model_works_with_infer_signature_multi_input_example(tmp_path):
|
|
model = SimpleChatModel()
|
|
params_subset = {
|
|
"max_tokens": 100,
|
|
}
|
|
input_example = {
|
|
"messages": [
|
|
{
|
|
"role": "assistant",
|
|
"content": "You are in helpful assistant!",
|
|
},
|
|
{
|
|
"role": "user",
|
|
"content": "What is Retrieval-augmented Generation?",
|
|
},
|
|
],
|
|
**params_subset,
|
|
}
|
|
with mlflow.start_run():
|
|
model_info = mlflow.pyfunc.log_model(
|
|
name="model", python_model=model, input_example=input_example
|
|
)
|
|
assert model_info.signature.inputs == CHAT_MODEL_INPUT_SCHEMA
|
|
assert model_info.signature.outputs == CHAT_MODEL_OUTPUT_SCHEMA
|
|
mlflow_model = Model.load(model_info.model_uri)
|
|
local_path = _download_artifact_from_uri(model_info.model_uri)
|
|
assert mlflow_model.load_input_example(local_path) == {
|
|
"messages": input_example["messages"],
|
|
**params_subset,
|
|
}
|
|
|
|
inference_payload = load_serving_example(model_info.model_uri)
|
|
response = pyfunc_serve_and_score_model(
|
|
model_uri=model_info.model_uri,
|
|
data=inference_payload,
|
|
content_type="application/json",
|
|
extra_args=["--env-manager", "local"],
|
|
)
|
|
|
|
expect_status_code(response, 200)
|
|
choices = json.loads(response.content)["choices"]
|
|
assert choices[0]["message"]["content"] == json.dumps(input_example["messages"])
|
|
assert json.loads(choices[1]["message"]["content"]) == {
|
|
**DEFAULT_PARAMS,
|
|
**params_subset,
|
|
}
|
|
|
|
|
|
def test_chat_model_predict_stream(tmp_path):
|
|
model = SimpleChatModel()
|
|
mlflow.pyfunc.save_model(python_model=model, path=tmp_path)
|
|
|
|
loaded_model = mlflow.pyfunc.load_model(tmp_path)
|
|
messages = [
|
|
{"role": "system", "content": "You are a helpful assistant"},
|
|
{"role": "user", "content": "Hello!"},
|
|
]
|
|
|
|
responses = list(loaded_model.predict_stream({"messages": messages}))
|
|
for i, resp in enumerate(responses[:-1]):
|
|
assert resp["choices"][0]["delta"]["content"] == f"message {i}"
|
|
|
|
assert responses[-1]["choices"][0]["delta"] == {}
|
|
|
|
|
|
def test_chat_model_can_receive_and_return_metadata():
|
|
messages = [{"role": "user", "content": "Hello!"}]
|
|
params = {
|
|
"custom_inputs": {"image_url": "example", "detail": "high", "other_dict": {"key": "value"}},
|
|
}
|
|
input_example = {
|
|
"messages": messages,
|
|
**params,
|
|
}
|
|
|
|
model = ChatModelWithMetadata()
|
|
with mlflow.start_run():
|
|
model_info = mlflow.pyfunc.log_model(
|
|
name="model",
|
|
python_model=model,
|
|
input_example=input_example,
|
|
)
|
|
|
|
loaded_model = mlflow.pyfunc.load_model(model_info.model_uri)
|
|
|
|
# test that it works for normal pyfunc predict
|
|
response = loaded_model.predict({"messages": messages, **params})
|
|
assert response["custom_outputs"] == params["custom_inputs"]
|
|
|
|
# test that it works in serving
|
|
inference_payload = load_serving_example(model_info.model_uri)
|
|
response = pyfunc_serve_and_score_model(
|
|
model_uri=model_info.model_uri,
|
|
data=inference_payload,
|
|
content_type="application/json",
|
|
extra_args=["--env-manager", "local"],
|
|
)
|
|
|
|
serving_response = json.loads(response.content)
|
|
assert serving_response["custom_outputs"] == params["custom_inputs"]
|
|
|
|
|
|
def test_chat_model_can_use_tool_calls():
|
|
messages = [{"role": "user", "content": "What's the weather?"}]
|
|
|
|
weather_tool = (
|
|
FunctionToolDefinition(
|
|
name="get_weather",
|
|
description="Get the weather for your current location",
|
|
parameters=ToolParamsSchema(
|
|
{
|
|
"city": {
|
|
"type": "string",
|
|
"description": "The city to get the weather for",
|
|
},
|
|
"unit": {"type": "string", "enum": ["F", "C"]},
|
|
},
|
|
required=["city", "unit"],
|
|
),
|
|
)
|
|
.to_tool_definition()
|
|
.to_dict()
|
|
)
|
|
|
|
example = {
|
|
"messages": messages,
|
|
"tools": [weather_tool],
|
|
}
|
|
|
|
model = ChatModelWithToolCalling()
|
|
with mlflow.start_run():
|
|
model_info = mlflow.pyfunc.log_model(
|
|
name="model",
|
|
python_model=model,
|
|
input_example=example,
|
|
)
|
|
|
|
loaded_model = mlflow.pyfunc.load_model(model_info.model_uri)
|
|
response = loaded_model.predict(example)
|
|
|
|
model_tool_calls = response["choices"][0]["message"]["tool_calls"]
|
|
assert json.loads(model_tool_calls[0]["function"]["arguments"]) == {
|
|
"city": "some_value",
|
|
"unit": "some_value",
|
|
}
|
|
|
|
|
|
def test_chat_model_without_context_in_predict():
|
|
response = ChatCompletionResponse(
|
|
choices=[ChatChoice(message=ChatMessage(role="assistant", content="hi"))]
|
|
)
|
|
chunk_response = ChatCompletionChunk(
|
|
choices=[ChatChunkChoice(delta=ChatChoiceDelta(role="assistant", content="hi"))]
|
|
)
|
|
|
|
class Model(mlflow.pyfunc.ChatModel):
|
|
def predict(self, messages: list[ChatMessage], params: ChatParams):
|
|
return response
|
|
|
|
def predict_stream(self, messages: list[ChatMessage], params: ChatParams):
|
|
yield chunk_response
|
|
|
|
model = Model()
|
|
messages = [ChatMessage(role="user", content="hello?", name="chat")]
|
|
assert model.predict(messages, ChatParams()) == response
|
|
assert next(iter(model.predict_stream(messages, ChatParams()))) == chunk_response
|
|
|
|
with mlflow.start_run():
|
|
model_info = mlflow.pyfunc.log_model(
|
|
name="model", python_model=model, input_example=messages
|
|
)
|
|
pyfunc_model = mlflow.pyfunc.load_model(model_info.model_uri)
|
|
input_data = {"messages": [{"role": "user", "content": "hello"}]}
|
|
assert pyfunc_model.predict(input_data) == response.to_dict()
|
|
assert next(iter(pyfunc_model.predict_stream(input_data))) == chunk_response.to_dict()
|