wehub-resource-sync
387 lines
13 KiB
Python
387 lines
13 KiB
Python
import json
|
|
from dataclasses import asdict, dataclass
|
|
|
|
import numpy as np
|
|
import pandas as pd
|
|
import pydantic
|
|
import pyspark
|
|
import pytest
|
|
from sklearn.ensemble import RandomForestRegressor
|
|
|
|
import mlflow
|
|
from mlflow.exceptions import MlflowException
|
|
from mlflow.models import Model, ModelSignature, infer_signature, rag_signatures, set_signature
|
|
from mlflow.models.model import get_model_info
|
|
from mlflow.types import DataType
|
|
from mlflow.types.schema import (
|
|
Array,
|
|
ColSpec,
|
|
ParamSchema,
|
|
ParamSpec,
|
|
Schema,
|
|
TensorSpec,
|
|
convert_dataclass_to_schema,
|
|
)
|
|
from mlflow.types.utils import InvalidDataForSignatureInferenceError
|
|
|
|
|
|
def test_model_signature_with_colspec():
|
|
signature1 = ModelSignature(
|
|
inputs=Schema([ColSpec(DataType.boolean), ColSpec(DataType.binary)]),
|
|
outputs=Schema([
|
|
ColSpec(name=None, type=DataType.double),
|
|
ColSpec(name=None, type=DataType.double),
|
|
]),
|
|
)
|
|
signature2 = ModelSignature(
|
|
inputs=Schema([ColSpec(DataType.boolean), ColSpec(DataType.binary)]),
|
|
outputs=Schema([
|
|
ColSpec(name=None, type=DataType.double),
|
|
ColSpec(name=None, type=DataType.double),
|
|
]),
|
|
)
|
|
assert signature1 == signature2
|
|
signature3 = ModelSignature(
|
|
inputs=Schema([ColSpec(DataType.boolean), ColSpec(DataType.binary)]),
|
|
outputs=Schema([
|
|
ColSpec(name=None, type=DataType.float),
|
|
ColSpec(name=None, type=DataType.double),
|
|
]),
|
|
)
|
|
assert signature3 != signature1
|
|
as_json = json.dumps(signature1.to_dict())
|
|
signature4 = ModelSignature.from_dict(json.loads(as_json))
|
|
assert signature1 == signature4
|
|
signature5 = ModelSignature(
|
|
inputs=Schema([ColSpec(DataType.boolean), ColSpec(DataType.binary)]), outputs=None
|
|
)
|
|
as_json = json.dumps(signature5.to_dict())
|
|
signature6 = ModelSignature.from_dict(json.loads(as_json))
|
|
assert signature5 == signature6
|
|
|
|
|
|
def test_model_signature_with_tensorspec():
|
|
signature1 = ModelSignature(
|
|
inputs=Schema([TensorSpec(np.dtype("float"), (-1, 28, 28))]),
|
|
outputs=Schema([TensorSpec(np.dtype("float"), (-1, 10))]),
|
|
)
|
|
signature2 = ModelSignature(
|
|
inputs=Schema([TensorSpec(np.dtype("float"), (-1, 28, 28))]),
|
|
outputs=Schema([TensorSpec(np.dtype("float"), (-1, 10))]),
|
|
)
|
|
# Single type mismatch
|
|
assert signature1 == signature2
|
|
signature3 = ModelSignature(
|
|
inputs=Schema([TensorSpec(np.dtype("float"), (-1, 28, 28))]),
|
|
outputs=Schema([TensorSpec(np.dtype("int"), (-1, 10))]),
|
|
)
|
|
assert signature3 != signature1
|
|
# Name mismatch
|
|
signature4 = ModelSignature(
|
|
inputs=Schema([TensorSpec(np.dtype("float"), (-1, 28, 28))]),
|
|
outputs=Schema([TensorSpec(np.dtype("float"), (-1, 10), "mismatch")]),
|
|
)
|
|
assert signature3 != signature4
|
|
as_json = json.dumps(signature1.to_dict())
|
|
signature5 = ModelSignature.from_dict(json.loads(as_json))
|
|
assert signature1 == signature5
|
|
|
|
# Test with name
|
|
signature6 = ModelSignature(
|
|
inputs=Schema([
|
|
TensorSpec(np.dtype("float"), (-1, 28, 28), name="image"),
|
|
TensorSpec(np.dtype("int"), (-1, 10), name="metadata"),
|
|
]),
|
|
outputs=Schema([TensorSpec(np.dtype("float"), (-1, 10), name="outputs")]),
|
|
)
|
|
signature7 = ModelSignature(
|
|
inputs=Schema([
|
|
TensorSpec(np.dtype("float"), (-1, 28, 28), name="image"),
|
|
TensorSpec(np.dtype("int"), (-1, 10), name="metadata"),
|
|
]),
|
|
outputs=Schema([TensorSpec(np.dtype("float"), (-1, 10), name="outputs")]),
|
|
)
|
|
assert signature6 == signature7
|
|
assert signature1 != signature6
|
|
|
|
# Test w/o output
|
|
signature8 = ModelSignature(
|
|
inputs=Schema([TensorSpec(np.dtype("float"), (-1, 28, 28))]), outputs=None
|
|
)
|
|
as_json = json.dumps(signature8.to_dict())
|
|
signature9 = ModelSignature.from_dict(json.loads(as_json))
|
|
assert signature8 == signature9
|
|
|
|
|
|
def test_model_signature_with_colspec_and_tensorspec():
|
|
signature1 = ModelSignature(inputs=Schema([ColSpec(DataType.double)]))
|
|
signature2 = ModelSignature(inputs=Schema([TensorSpec(np.dtype("float"), (-1, 28, 28))]))
|
|
assert signature1 != signature2
|
|
assert signature2 != signature1
|
|
|
|
signature3 = ModelSignature(
|
|
inputs=Schema([ColSpec(DataType.double)]),
|
|
outputs=Schema([TensorSpec(np.dtype("float"), (-1, 28, 28))]),
|
|
)
|
|
signature4 = ModelSignature(
|
|
inputs=Schema([ColSpec(DataType.double)]),
|
|
outputs=Schema([ColSpec(DataType.double)]),
|
|
)
|
|
assert signature3 != signature4
|
|
assert signature4 != signature3
|
|
|
|
|
|
def test_signature_inference_infers_input_and_output_as_expected():
|
|
sig0 = infer_signature(np.array([1]))
|
|
assert sig0.inputs is not None
|
|
assert sig0.outputs is None
|
|
sig1 = infer_signature(np.array([1]), np.array([1]))
|
|
assert sig1.inputs == sig0.inputs
|
|
assert sig1.outputs == sig0.inputs
|
|
|
|
|
|
def test_infer_signature_on_nested_array():
|
|
signature = infer_signature(
|
|
model_input=[{"queries": [["a", "b", "c"], ["d", "e"], []]}],
|
|
model_output=[{"answers": [["f", "g"], ["h"]]}],
|
|
)
|
|
assert signature.inputs == Schema([ColSpec(Array(Array(DataType.string)), name="queries")])
|
|
assert signature.outputs == Schema([ColSpec(Array(Array(DataType.string)), name="answers")])
|
|
|
|
signature = infer_signature(
|
|
model_input=[
|
|
{
|
|
"inputs": [
|
|
np.array([["a", "b"], ["c", "d"]]),
|
|
np.array([["e", "f"], ["g", "h"]]),
|
|
]
|
|
}
|
|
],
|
|
model_output=[{"outputs": [np.int32(5), np.int32(6)]}],
|
|
)
|
|
assert signature.inputs == Schema([
|
|
ColSpec(Array(Array(Array(DataType.string))), name="inputs")
|
|
])
|
|
assert signature.outputs == Schema([ColSpec(Array(DataType.integer), name="outputs")])
|
|
|
|
|
|
def test_infer_signature_on_list_of_dictionaries():
|
|
signature = infer_signature(
|
|
model_input=[{"query": "test query"}],
|
|
model_output=[
|
|
{
|
|
"output": "Output from the LLM",
|
|
"candidate_ids": ["412", "1233"],
|
|
"candidate_sources": ["file1.md", "file201.md"],
|
|
}
|
|
],
|
|
)
|
|
assert signature.inputs == Schema([ColSpec(DataType.string, name="query")])
|
|
assert signature.outputs == Schema([
|
|
ColSpec(DataType.string, name="output"),
|
|
ColSpec(Array(DataType.string), name="candidate_ids"),
|
|
ColSpec(Array(DataType.string), name="candidate_sources"),
|
|
])
|
|
|
|
|
|
def test_signature_inference_infers_datime_types_as_expected():
|
|
col_name = "datetime_col"
|
|
test_datetime = np.datetime64("2021-01-01")
|
|
test_series = pd.Series(pd.to_datetime([test_datetime]))
|
|
test_df = test_series.to_frame(col_name)
|
|
|
|
signature = infer_signature(test_series)
|
|
assert signature.inputs == Schema([ColSpec(DataType.datetime)])
|
|
|
|
signature = infer_signature(test_df)
|
|
assert signature.inputs == Schema([ColSpec(DataType.datetime, name=col_name)])
|
|
|
|
with pyspark.sql.SparkSession.builder.getOrCreate() as spark:
|
|
spark_df = spark.range(1).selectExpr(
|
|
"current_timestamp() as timestamp", "current_date() as date"
|
|
)
|
|
signature = infer_signature(spark_df)
|
|
assert signature.inputs == Schema([
|
|
ColSpec(DataType.datetime, name="timestamp"),
|
|
ColSpec(DataType.datetime, name="date"),
|
|
])
|
|
|
|
|
|
def test_set_signature_to_logged_model():
|
|
artifact_path = "regr-model"
|
|
with mlflow.start_run():
|
|
model_info = mlflow.sklearn.log_model(RandomForestRegressor(), name=artifact_path)
|
|
signature = infer_signature(np.array([1]))
|
|
set_signature(model_info.model_uri, signature)
|
|
model_info = get_model_info(model_info.model_uri)
|
|
assert model_info.signature == signature
|
|
|
|
|
|
def test_set_signature_to_saved_model(tmp_path):
|
|
model_path = str(tmp_path)
|
|
mlflow.sklearn.save_model(
|
|
RandomForestRegressor(),
|
|
model_path,
|
|
serialization_format=mlflow.sklearn.SERIALIZATION_FORMAT_CLOUDPICKLE,
|
|
)
|
|
signature = infer_signature(np.array([1]))
|
|
set_signature(model_path, signature)
|
|
assert Model.load(model_path).signature == signature
|
|
|
|
|
|
def test_set_signature_overwrite():
|
|
artifact_path = "regr-model"
|
|
with mlflow.start_run():
|
|
model_info = mlflow.sklearn.log_model(
|
|
RandomForestRegressor(),
|
|
name=artifact_path,
|
|
signature=infer_signature(np.array([1])),
|
|
)
|
|
new_signature = infer_signature(np.array([1]), np.array([1]))
|
|
set_signature(model_info.model_uri, new_signature)
|
|
model_info = get_model_info(model_info.model_uri)
|
|
assert model_info.signature == new_signature
|
|
|
|
|
|
def test_cannot_set_signature_on_models_scheme_uris():
|
|
signature = infer_signature(np.array([1]))
|
|
with pytest.raises(
|
|
MlflowException,
|
|
match="Model URIs with the `models:/<name>/<version>` scheme are not supported.",
|
|
):
|
|
set_signature("models:/dummy_model@champion", signature)
|
|
|
|
|
|
def test_signature_construction():
|
|
signature = ModelSignature(inputs=Schema([ColSpec(DataType.binary)]))
|
|
assert signature.to_dict() == {
|
|
"inputs": '[{"type": "binary", "required": true}]',
|
|
"outputs": None,
|
|
"params": None,
|
|
}
|
|
|
|
signature = ModelSignature(outputs=Schema([ColSpec(DataType.double)]))
|
|
assert signature.to_dict() == {
|
|
"inputs": None,
|
|
"outputs": '[{"type": "double", "required": true}]',
|
|
"params": None,
|
|
}
|
|
|
|
signature = ModelSignature(params=ParamSchema([ParamSpec("param1", DataType.string, "test")]))
|
|
assert signature.to_dict() == {
|
|
"inputs": None,
|
|
"outputs": None,
|
|
"params": '[{"name": "param1", "default": "test", "shape": null, "type": "string"}]',
|
|
}
|
|
|
|
|
|
def test_signature_with_errors():
|
|
with pytest.raises(
|
|
TypeError,
|
|
match=r"inputs must be either None, mlflow.models.signature.Schema, or a dataclass",
|
|
):
|
|
ModelSignature(inputs=1)
|
|
|
|
with pytest.raises(
|
|
ValueError, match=r"At least one of inputs, outputs or params must be provided"
|
|
):
|
|
ModelSignature()
|
|
|
|
|
|
def test_signature_for_rag():
|
|
signature = ModelSignature(
|
|
inputs=rag_signatures.ChatCompletionRequest(),
|
|
outputs=rag_signatures.ChatCompletionResponse(),
|
|
)
|
|
signature_dict = signature.to_dict()
|
|
assert signature_dict == {
|
|
"inputs": (
|
|
'[{"type": "array", "items": {"type": "object", "properties": '
|
|
'{"content": {"type": "string", "required": true}, '
|
|
'"role": {"type": "string", "required": true}}}, '
|
|
'"name": "messages", "required": true}]'
|
|
),
|
|
"outputs": (
|
|
'[{"type": "array", "items": {"type": "object", "properties": '
|
|
'{"finish_reason": {"type": "string", "required": true}, '
|
|
'"index": {"type": "long", "required": true}, '
|
|
'"message": {"type": "object", "properties": '
|
|
'{"content": {"type": "string", "required": true}, '
|
|
'"role": {"type": "string", "required": true}}, '
|
|
'"required": true}}}, "name": "choices", "required": true}, '
|
|
'{"type": "string", "name": "object", "required": true}]'
|
|
),
|
|
"params": None,
|
|
}
|
|
|
|
|
|
def test_infer_signature_and_convert_dataclass_to_schema_for_rag():
|
|
inferred_signature = infer_signature(
|
|
asdict(rag_signatures.ChatCompletionRequest()),
|
|
asdict(rag_signatures.ChatCompletionResponse()),
|
|
)
|
|
input_schema = convert_dataclass_to_schema(rag_signatures.ChatCompletionRequest())
|
|
output_schema = convert_dataclass_to_schema(rag_signatures.ChatCompletionResponse())
|
|
assert inferred_signature.inputs == input_schema
|
|
assert inferred_signature.outputs == output_schema
|
|
|
|
|
|
def test_infer_signature_with_dataclass():
|
|
inferred_signature = infer_signature(
|
|
rag_signatures.ChatCompletionRequest(),
|
|
rag_signatures.ChatCompletionResponse(),
|
|
)
|
|
input_schema = convert_dataclass_to_schema(rag_signatures.ChatCompletionRequest())
|
|
output_schema = convert_dataclass_to_schema(rag_signatures.ChatCompletionResponse())
|
|
assert inferred_signature.inputs == input_schema
|
|
assert inferred_signature.outputs == output_schema
|
|
|
|
|
|
@dataclass
|
|
class CustomInput:
|
|
id: int = 0
|
|
|
|
|
|
@dataclass
|
|
class CustomOutput:
|
|
id: int = 0
|
|
|
|
|
|
@dataclass
|
|
class FlexibleChatCompletionRequest(rag_signatures.ChatCompletionRequest):
|
|
custom_input: CustomInput | None = None
|
|
|
|
|
|
@dataclass
|
|
class FlexibleChatCompletionResponse(rag_signatures.ChatCompletionResponse):
|
|
custom_output: CustomOutput | None = None
|
|
|
|
|
|
def test_infer_signature_with_optional_and_child_dataclass():
|
|
inferred_signature = infer_signature(
|
|
FlexibleChatCompletionRequest(),
|
|
FlexibleChatCompletionResponse(),
|
|
)
|
|
custom_input_schema = next(
|
|
schema for schema in inferred_signature.inputs.to_dict() if schema["name"] == "custom_input"
|
|
)
|
|
assert custom_input_schema["required"] is False
|
|
assert "id" in custom_input_schema["properties"]
|
|
assert any(
|
|
schema for schema in inferred_signature.inputs.to_dict() if schema["name"] == "messages"
|
|
)
|
|
|
|
|
|
def test_infer_signature_for_pydantic_objects_error():
|
|
class Message(pydantic.BaseModel):
|
|
content: str
|
|
role: str
|
|
|
|
m = Message(content="test", role="user")
|
|
with pytest.raises(
|
|
InvalidDataForSignatureInferenceError,
|
|
match=r"MLflow does not support inferring model signature from "
|
|
r"input example with Pydantic objects",
|
|
):
|
|
infer_signature([m])
|