wehub-resource-sync
104 lines
3.5 KiB
Python
104 lines
3.5 KiB
Python
import json
|
|
import os
|
|
|
|
import pytest
|
|
|
|
from mlflow.entities.param import Param
|
|
from mlflow.entities.run_status import RunStatus
|
|
from mlflow.store.artifact.artifact_repository_registry import get_artifact_repository
|
|
from mlflow.tracking._tracking_service.utils import _get_store
|
|
from mlflow.utils.promptlab_utils import (
|
|
_create_promptlab_run_impl,
|
|
create_eval_results_json,
|
|
)
|
|
|
|
prompt_parameters = [
|
|
Param(key="question", value="my_question"),
|
|
Param(key="context", value="my_context"),
|
|
]
|
|
model_input = "answer this question: my_question using the following context: my_context"
|
|
model_output = "my_answer"
|
|
model_output_parameters = [
|
|
Param(key="tokens", value="10"),
|
|
Param(key="latency", value="100"),
|
|
]
|
|
|
|
|
|
def test_eval_results_file():
|
|
eval_results_file = create_eval_results_json(
|
|
prompt_parameters, model_input, model_output_parameters, model_output
|
|
)
|
|
expected_eval_results_json = {
|
|
"columns": ["question", "context", "prompt", "output", "tokens", "latency"],
|
|
"data": [
|
|
[
|
|
"my_question",
|
|
"my_context",
|
|
"answer this question: my_question using the following context: my_context",
|
|
"my_answer",
|
|
"10",
|
|
"100",
|
|
]
|
|
],
|
|
}
|
|
assert json.loads(eval_results_file) == expected_eval_results_json
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
"MLFLOW_SKINNY" in os.environ,
|
|
reason="Skinny does not support the np or pandas dependencies",
|
|
)
|
|
def test_create_promptlab_run(db_uri):
|
|
store = _get_store(db_uri)
|
|
exp_id = store.create_experiment("test_create_promptlab_run")
|
|
run = _create_promptlab_run_impl(
|
|
store,
|
|
experiment_id=exp_id,
|
|
run_name="my_promptlab_run",
|
|
tags=[],
|
|
prompt_template="my_prompt_template",
|
|
prompt_parameters=[Param("prompt_param_key", "prompt_param_value")],
|
|
model_route="my_route",
|
|
model_parameters=[Param("temperature", "0.1")],
|
|
model_input="",
|
|
model_output_parameters=[Param("output_param_key", "output_param_value")],
|
|
model_output="my_output",
|
|
mlflow_version="1.0.0",
|
|
user_id="user",
|
|
start_time=1,
|
|
)
|
|
assert run.info.run_id is not None
|
|
assert run.info.status == RunStatus.to_string(RunStatus.FINISHED)
|
|
|
|
assert run.data.params["prompt_template"] == "my_prompt_template"
|
|
assert run.data.params["model_route"] == "my_route"
|
|
assert run.data.params["temperature"] == "0.1"
|
|
|
|
assert run.data.tags["mlflow.runName"] == "my_promptlab_run"
|
|
assert (
|
|
run.data.tags["mlflow.loggedArtifacts"]
|
|
== '[{"path": "eval_results_table.json", "type": "table"}]'
|
|
)
|
|
assert run.data.tags["mlflow.runSourceType"] == "PROMPT_ENGINEERING"
|
|
assert run.data.tags["mlflow.log-model.history"] is not None
|
|
|
|
# list the files in the model folder
|
|
artifact_location = run.info.artifact_uri
|
|
artifact_repo = get_artifact_repository(artifact_location)
|
|
|
|
artifact_files = [f.path for f in artifact_repo.list_artifacts()]
|
|
assert "eval_results_table.json" in artifact_files
|
|
assert "model" in artifact_files
|
|
|
|
model_files = [f.path for f in artifact_repo.list_artifacts("model")]
|
|
assert "model/MLmodel" in model_files
|
|
assert "model/python_env.yaml" in model_files
|
|
assert "model/conda.yaml" in model_files
|
|
assert "model/requirements.txt" in model_files
|
|
assert "model/input_example.json" in model_files
|
|
|
|
# try to load the model
|
|
import mlflow.pyfunc
|
|
|
|
mlflow.pyfunc.load_model(f"{artifact_location}/model")
|