wehub-resource-sync
93 lines
3.6 KiB
Python
93 lines
3.6 KiB
Python
import os
|
|
from unittest import mock
|
|
|
|
import pytest
|
|
|
|
from mlflow.exceptions import MlflowException
|
|
from mlflow.models import ModelConfig
|
|
|
|
dir_path = os.path.dirname(os.path.abspath(__file__))
|
|
VALID_CONFIG_PATH = os.path.join(dir_path, "configs/config.yaml")
|
|
VALID_CONFIG_PATH_2 = os.path.join(dir_path, "configs/config_2.yaml")
|
|
|
|
|
|
def test_config_not_set():
|
|
with pytest.raises(
|
|
FileNotFoundError, match="Config file is not provided which is needed to load the model."
|
|
):
|
|
ModelConfig()
|
|
|
|
|
|
def test_config_not_found():
|
|
with pytest.raises(FileNotFoundError, match="Config file 'nonexistent.yaml' not found."):
|
|
ModelConfig(development_config="nonexistent.yaml")
|
|
|
|
|
|
def test_config_invalid_yaml(tmp_path):
|
|
tmp_file = tmp_path / "invalid_config.yaml"
|
|
tmp_file.write_text("invalid_yaml: \n - this is not valid \n-yaml")
|
|
config = ModelConfig(development_config=str(tmp_file))
|
|
with pytest.raises(MlflowException, match="Error parsing YAML file: "):
|
|
config.get("key")
|
|
|
|
|
|
def test_config_key_not_found():
|
|
config = ModelConfig(development_config=VALID_CONFIG_PATH)
|
|
with pytest.raises(KeyError, match="Key 'key' not found in configuration: "):
|
|
config.get("key")
|
|
|
|
|
|
def test_config_setup_correctly():
|
|
config = ModelConfig(development_config=VALID_CONFIG_PATH)
|
|
assert config.get("llm_parameters").get("temperature") == 0.01
|
|
|
|
|
|
def test_config_setup_correctly_with_mlflow_langchain():
|
|
with mock.patch("mlflow.models.model_config.__mlflow_model_config__", new=VALID_CONFIG_PATH):
|
|
config = ModelConfig(development_config="nonexistent.yaml")
|
|
assert config.get("llm_parameters").get("temperature") == 0.01
|
|
|
|
|
|
def test_config_setup_with_mlflow_langchain_path():
|
|
with mock.patch("mlflow.models.model_config.__mlflow_model_config__", new=VALID_CONFIG_PATH_2):
|
|
# here the config.yaml has the max_tokens set to 500
|
|
# where as the config_2.yaml has it set to 200.
|
|
# Here we give preference to the __mlflow_model_config__.
|
|
config = ModelConfig(development_config=VALID_CONFIG_PATH)
|
|
assert config.get("llm_parameters").get("max_tokens") == 200
|
|
|
|
|
|
def test_config_development_config_must_be_specified_with_keyword():
|
|
with pytest.raises(TypeError, match="1 positional argument but 2 were given"):
|
|
ModelConfig(VALID_CONFIG_PATH_2)
|
|
|
|
|
|
def test_config_development_config_is_a_dict():
|
|
config = ModelConfig(development_config={"llm_parameters": {"temperature": 0.01}})
|
|
assert config.get("llm_parameters").get("temperature") == 0.01
|
|
|
|
|
|
def test_config_setup_correctly_errors_with_no_config_path():
|
|
with mock.patch("mlflow.models.model_config.__mlflow_model_config__", new=""):
|
|
with pytest.raises(
|
|
FileNotFoundError,
|
|
match="Config file is not provided which is needed to load the model.",
|
|
):
|
|
ModelConfig(development_config=VALID_CONFIG_PATH)
|
|
|
|
|
|
def test_config_development_config_to_dict():
|
|
config = ModelConfig(development_config={"llm_parameters": {"temperature": 0.01}})
|
|
assert config.to_dict() == {"llm_parameters": {"temperature": 0.01}}
|
|
|
|
config = ModelConfig(development_config=VALID_CONFIG_PATH)
|
|
assert config.to_dict() == {
|
|
"embedding_model_query_instructions": "Represent this sentence for searching "
|
|
"relevant passages:",
|
|
"llm_model": "databricks-dbrx-instruct",
|
|
"llm_prompt_template": "You are a trustful assistant.",
|
|
"retriever_config": {"k": 5, "use_mmr": False},
|
|
"llm_parameters": {"temperature": 0.01, "max_tokens": 500},
|
|
"llm_prompt_template_variables": ["chat_history", "context", "question"],
|
|
}
|