wehub-resource-sync
60 lines
2.1 KiB
Python
60 lines
2.1 KiB
Python
import json
|
|
from dataclasses import asdict
|
|
|
|
import mlflow
|
|
from mlflow.models.model import Model
|
|
from mlflow.models.rag_signatures import (
|
|
ChainCompletionChoice,
|
|
ChatCompletionRequest,
|
|
ChatCompletionResponse,
|
|
Message,
|
|
)
|
|
from mlflow.models.signature import ModelSignature
|
|
|
|
from tests.helper_functions import expect_status_code, pyfunc_serve_and_score_model
|
|
|
|
|
|
class TestRagModel(mlflow.pyfunc.PythonModel):
|
|
def predict(self, context, model_input: ChatCompletionRequest):
|
|
message = model_input.messages[0].content
|
|
# return the message back
|
|
return asdict(
|
|
ChatCompletionResponse(
|
|
choices=[ChainCompletionChoice(message=Message(role="assistant", content=message))]
|
|
# NB: intentionally validating the default population of the object field
|
|
)
|
|
)
|
|
|
|
|
|
def test_rag_model_works_with_type_hint(tmp_path):
|
|
model = TestRagModel()
|
|
signature = ModelSignature(inputs=ChatCompletionRequest(), outputs=ChatCompletionResponse())
|
|
input_example = {"messages": [{"role": "user", "content": "What is mlflow?"}]}
|
|
mlflow.pyfunc.save_model(
|
|
python_model=model, path=tmp_path, signature=signature, input_example=input_example
|
|
)
|
|
|
|
# test that the model can be loaded and invoked
|
|
loaded_model = mlflow.pyfunc.load_model(tmp_path)
|
|
|
|
response = loaded_model.predict(input_example)
|
|
assert response["choices"][0]["message"]["content"] == "What is mlflow?"
|
|
assert response["object"] == "chat.completion"
|
|
|
|
# confirm the input example is set
|
|
mlflow_model = Model.load(tmp_path)
|
|
assert mlflow_model.load_input_example(tmp_path) == input_example
|
|
|
|
# test that the model can be served
|
|
response = pyfunc_serve_and_score_model(
|
|
model_uri=tmp_path,
|
|
data=json.dumps(input_example),
|
|
content_type="application/json",
|
|
extra_args=["--env-manager", "local"],
|
|
)
|
|
|
|
expect_status_code(response, 200)
|
|
json_response = json.loads(response.content)
|
|
assert json_response["choices"][0]["message"]["content"] == "What is mlflow?"
|
|
assert json_response["object"] == "chat.completion"
|