wehub-resource-sync
479 lines
17 KiB
Python
479 lines
17 KiB
Python
from typing import Any
|
|
|
|
import pandas as pd
|
|
import pytest
|
|
|
|
import mlflow
|
|
from mlflow.entities.model_registry import PromptModelConfig
|
|
from mlflow.exceptions import MlflowException
|
|
from mlflow.genai.datasets import create_dataset
|
|
from mlflow.genai.optimize.optimize import optimize_prompts
|
|
from mlflow.genai.optimize.optimizers.base import BasePromptOptimizer
|
|
from mlflow.genai.optimize.types import EvaluationResultRecord, PromptOptimizerOutput
|
|
from mlflow.genai.prompts import register_prompt
|
|
from mlflow.genai.scorers import scorer
|
|
from mlflow.models.model import PromptVersion
|
|
from mlflow.utils.import_hooks import _post_import_hooks
|
|
|
|
|
|
class MockPromptOptimizer(BasePromptOptimizer):
|
|
def __init__(self, reflection_model="openai:/gpt-4o-mini"):
|
|
self.model_name = reflection_model
|
|
|
|
def optimize(
|
|
self,
|
|
eval_fn: Any,
|
|
train_data: list[dict[str, Any]],
|
|
target_prompts: dict[str, str],
|
|
enable_tracking: bool = True,
|
|
) -> PromptOptimizerOutput:
|
|
optimized_prompts = {}
|
|
for prompt_name, template in target_prompts.items():
|
|
# Simple optimization: add "Be precise and accurate. " prefix
|
|
optimized_prompts[prompt_name] = f"Be precise and accurate. {template}"
|
|
|
|
# Verify the optimization by calling eval_fn (only if provided)
|
|
if eval_fn is not None:
|
|
eval_fn(optimized_prompts, train_data)
|
|
|
|
return PromptOptimizerOutput(
|
|
optimized_prompts=optimized_prompts,
|
|
initial_eval_score=0.5,
|
|
final_eval_score=0.9,
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def sample_translation_prompt() -> PromptVersion:
|
|
return register_prompt(
|
|
name="test_translation_prompt",
|
|
template="Translate the following text to {{language}}: {{input_text}}",
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def sample_summarization_prompt() -> PromptVersion:
|
|
return register_prompt(
|
|
name="test_summarization_prompt",
|
|
template="Summarize this text: {{text}}",
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def sample_dataset() -> pd.DataFrame:
|
|
return pd.DataFrame({
|
|
"inputs": [
|
|
{"input_text": "Hello", "language": "Spanish"},
|
|
{"input_text": "World", "language": "French"},
|
|
{"input_text": "Goodbye", "language": "Spanish"},
|
|
],
|
|
"outputs": [
|
|
"Hola",
|
|
"Monde",
|
|
"Adiós",
|
|
],
|
|
})
|
|
|
|
|
|
@pytest.fixture
|
|
def sample_summarization_dataset() -> list[dict[str, Any]]:
|
|
return [
|
|
{
|
|
"inputs": {
|
|
"text": "This is a long document that needs to be summarized into key points."
|
|
},
|
|
"outputs": "Key points summary",
|
|
},
|
|
{
|
|
"inputs": {"text": "Another document with important information for summarization."},
|
|
"outputs": "Important info summary",
|
|
},
|
|
]
|
|
|
|
|
|
def sample_predict_fn(input_text: str, language: str) -> str:
|
|
mlflow.genai.load_prompt("prompts:/test_translation_prompt/1")
|
|
translations = {
|
|
("Hello", "Spanish"): "Hola",
|
|
("World", "French"): "Monde",
|
|
("Goodbye", "Spanish"): "Adiós",
|
|
}
|
|
|
|
# Verify that auto logging is enabled during the evaluation.
|
|
assert len(_post_import_hooks) > 0
|
|
return translations.get((input_text, language), f"translated_{input_text}")
|
|
|
|
|
|
def sample_summarization_fn(text: str) -> str:
|
|
return f"Summary of: {text[:20]}..."
|
|
|
|
|
|
@mlflow.genai.scorers.scorer(name="equivalence")
|
|
def equivalence(outputs, expectations):
|
|
return 1.0 if outputs == expectations["expected_response"] else 0.0
|
|
|
|
|
|
def test_optimize_prompts_single_prompt(
|
|
sample_translation_prompt: PromptVersion, sample_dataset: pd.DataFrame
|
|
):
|
|
mock_optimizer = MockPromptOptimizer()
|
|
|
|
result = optimize_prompts(
|
|
predict_fn=sample_predict_fn,
|
|
train_data=sample_dataset,
|
|
prompt_uris=[
|
|
f"prompts:/{sample_translation_prompt.name}/{sample_translation_prompt.version}"
|
|
],
|
|
optimizer=mock_optimizer,
|
|
scorers=[equivalence],
|
|
)
|
|
|
|
assert len(result.optimized_prompts) == 1
|
|
optimized_prompt = result.optimized_prompts[0]
|
|
assert optimized_prompt.name == sample_translation_prompt.name
|
|
assert optimized_prompt.version == sample_translation_prompt.version + 1
|
|
assert "Be precise and accurate." in optimized_prompt.template
|
|
expected_template = "Translate the following text to {{language}}: {{input_text}}"
|
|
assert expected_template in optimized_prompt.template
|
|
assert result.initial_eval_score == 0.5
|
|
assert result.final_eval_score == 0.9
|
|
|
|
|
|
def test_optimize_prompts_multiple_prompts(
|
|
sample_translation_prompt: PromptVersion,
|
|
sample_summarization_prompt: PromptVersion,
|
|
sample_dataset: pd.DataFrame,
|
|
):
|
|
mock_optimizer = MockPromptOptimizer()
|
|
|
|
result = optimize_prompts(
|
|
predict_fn=sample_predict_fn,
|
|
train_data=sample_dataset,
|
|
prompt_uris=[
|
|
f"prompts:/{sample_translation_prompt.name}/{sample_translation_prompt.version}",
|
|
f"prompts:/{sample_summarization_prompt.name}/{sample_summarization_prompt.version}",
|
|
],
|
|
optimizer=mock_optimizer,
|
|
scorers=[equivalence],
|
|
)
|
|
|
|
assert len(result.optimized_prompts) == 2
|
|
prompt_names = {prompt.name for prompt in result.optimized_prompts}
|
|
assert sample_translation_prompt.name in prompt_names
|
|
assert sample_summarization_prompt.name in prompt_names
|
|
assert result.initial_eval_score == 0.5
|
|
assert result.final_eval_score == 0.9
|
|
|
|
for prompt in result.optimized_prompts:
|
|
assert "Be precise and accurate." in prompt.template
|
|
|
|
|
|
def test_optimize_prompts_eval_function_behavior(
|
|
sample_translation_prompt: PromptVersion, sample_dataset: pd.DataFrame
|
|
):
|
|
class TestingOptimizer(BasePromptOptimizer):
|
|
def __init__(self):
|
|
self.model_name = "openai:/gpt-4o-mini"
|
|
self.eval_fn_calls = []
|
|
|
|
def optimize(self, eval_fn, dataset, target_prompts, enable_tracking=True):
|
|
# Test that eval_fn works correctly
|
|
test_prompts = {
|
|
"test_translation_prompt": "Prompt Candidate: "
|
|
"Translate {{input_text}} to {{language}}"
|
|
}
|
|
results = eval_fn(test_prompts, dataset)
|
|
self.eval_fn_calls.append((test_prompts, results))
|
|
|
|
# Verify results structure
|
|
assert isinstance(results, list)
|
|
assert len(results) == len(dataset)
|
|
for i, result in enumerate(results):
|
|
assert isinstance(result, EvaluationResultRecord)
|
|
assert result.inputs == dataset[i]["inputs"]
|
|
assert result.outputs == dataset[i]["outputs"]
|
|
assert result.score == 1
|
|
assert result.trace is not None
|
|
|
|
return PromptOptimizerOutput(optimized_prompts=target_prompts)
|
|
|
|
predict_called_count = 0
|
|
|
|
def predict_fn(input_text, language):
|
|
prompt = mlflow.genai.load_prompt("prompts:/test_translation_prompt/1").format(
|
|
input_text=input_text, language=language
|
|
)
|
|
nonlocal predict_called_count
|
|
# the first call to the predict_fn is the model check
|
|
if predict_called_count > 0:
|
|
# validate the prompt is replaced with the candidate prompt
|
|
assert "Prompt Candidate" in prompt
|
|
predict_called_count += 1
|
|
|
|
return sample_predict_fn(input_text=input_text, language=language)
|
|
|
|
testing_optimizer = TestingOptimizer()
|
|
|
|
optimize_prompts(
|
|
predict_fn=predict_fn,
|
|
train_data=sample_dataset,
|
|
prompt_uris=[
|
|
f"prompts:/{sample_translation_prompt.name}/{sample_translation_prompt.version}"
|
|
],
|
|
optimizer=testing_optimizer,
|
|
scorers=[equivalence],
|
|
)
|
|
|
|
assert len(testing_optimizer.eval_fn_calls) == 1
|
|
_, eval_results = testing_optimizer.eval_fn_calls[0]
|
|
assert len(eval_results) == 3 # Number of records in sample_dataset
|
|
assert predict_called_count == 4 # 3 records in sample_dataset + 1 for the prediction check
|
|
|
|
|
|
def test_optimize_prompts_with_list_dataset(
|
|
sample_translation_prompt: PromptVersion, sample_summarization_dataset: list[dict[str, Any]]
|
|
):
|
|
mock_optimizer = MockPromptOptimizer()
|
|
|
|
def summarization_predict_fn(text):
|
|
return f"Summary: {text[:10]}..."
|
|
|
|
result = optimize_prompts(
|
|
predict_fn=summarization_predict_fn,
|
|
train_data=sample_summarization_dataset,
|
|
prompt_uris=[
|
|
f"prompts:/{sample_translation_prompt.name}/{sample_translation_prompt.version}"
|
|
],
|
|
optimizer=mock_optimizer,
|
|
scorers=[equivalence],
|
|
)
|
|
|
|
assert len(result.optimized_prompts) == 1
|
|
assert result.initial_eval_score == 0.5
|
|
assert result.final_eval_score == 0.9
|
|
|
|
|
|
def test_optimize_prompts_with_model_name(
|
|
sample_translation_prompt: PromptVersion, sample_dataset: pd.DataFrame
|
|
):
|
|
class TestOptimizer(BasePromptOptimizer):
|
|
def __init__(self):
|
|
self.model_name = "test/custom-model"
|
|
|
|
def optimize(self, eval_fn, dataset, target_prompts, enable_tracking=True):
|
|
return PromptOptimizerOutput(optimized_prompts=target_prompts)
|
|
|
|
testing_optimizer = TestOptimizer()
|
|
|
|
result = optimize_prompts(
|
|
predict_fn=sample_predict_fn,
|
|
train_data=sample_dataset,
|
|
prompt_uris=[
|
|
f"prompts:/{sample_translation_prompt.name}/{sample_translation_prompt.version}"
|
|
],
|
|
optimizer=testing_optimizer,
|
|
scorers=[equivalence],
|
|
)
|
|
|
|
assert len(result.optimized_prompts) == 1
|
|
|
|
|
|
def test_optimize_prompts_warns_on_unused_prompt(
|
|
sample_translation_prompt: PromptVersion,
|
|
sample_summarization_prompt: PromptVersion,
|
|
sample_dataset: pd.DataFrame,
|
|
capsys,
|
|
):
|
|
mock_optimizer = MockPromptOptimizer()
|
|
|
|
# Create predict_fn that only uses translation prompt, not summarization prompt
|
|
def predict_fn_single_prompt(input_text, language):
|
|
prompt = mlflow.genai.load_prompt("prompts:/test_translation_prompt/1")
|
|
prompt.format(input_text=input_text, language=language)
|
|
return sample_predict_fn(input_text=input_text, language=language)
|
|
|
|
result = optimize_prompts(
|
|
predict_fn=predict_fn_single_prompt,
|
|
train_data=sample_dataset,
|
|
prompt_uris=[
|
|
f"prompts:/{sample_translation_prompt.name}/{sample_translation_prompt.version}",
|
|
f"prompts:/{sample_summarization_prompt.name}/{sample_summarization_prompt.version}",
|
|
],
|
|
optimizer=mock_optimizer,
|
|
scorers=[equivalence],
|
|
)
|
|
|
|
assert len(result.optimized_prompts) == 2
|
|
|
|
captured = capsys.readouterr()
|
|
assert "prompts were not used during evaluation" in captured.err
|
|
assert "test_summarization_prompt" in captured.err
|
|
|
|
|
|
def test_optimize_prompts_with_custom_scorers(
|
|
sample_translation_prompt: PromptVersion, sample_dataset: pd.DataFrame
|
|
):
|
|
# Create a custom scorer for case-insensitive matching
|
|
@scorer(name="case_insensitive_match")
|
|
def case_insensitive_match(outputs, expectations):
|
|
# Extract expected_response if expectations is a dict
|
|
if isinstance(expectations, dict) and "expected_response" in expectations:
|
|
expected_value = expectations["expected_response"]
|
|
else:
|
|
expected_value = expectations
|
|
return 1.0 if str(outputs).lower() == str(expected_value).lower() else 0.5
|
|
|
|
class MetricTestOptimizer(BasePromptOptimizer):
|
|
def __init__(self):
|
|
self.model_name = "openai:/gpt-4o-mini"
|
|
self.captured_scores = []
|
|
|
|
def optimize(self, eval_fn, dataset, target_prompts, enable_tracking=True):
|
|
# Run eval_fn and capture the scores
|
|
results = eval_fn(target_prompts, dataset)
|
|
self.captured_scores = [r.score for r in results]
|
|
return PromptOptimizerOutput(optimized_prompts=target_prompts)
|
|
|
|
testing_optimizer = MetricTestOptimizer()
|
|
|
|
# Create dataset with outputs that will test custom scorer
|
|
test_dataset = pd.DataFrame({
|
|
"inputs": [
|
|
{"input_text": "Hello", "language": "Spanish"},
|
|
{"input_text": "World", "language": "French"},
|
|
],
|
|
"outputs": ["HOLA", "monde"], # Different cases to test custom scorer
|
|
})
|
|
|
|
def predict_fn(input_text, language):
|
|
mlflow.genai.load_prompt("prompts:/test_translation_prompt/1")
|
|
# Return lowercase outputs
|
|
return {"Hello": "hola", "World": "monde"}.get(input_text, "unknown")
|
|
|
|
result = optimize_prompts(
|
|
predict_fn=predict_fn,
|
|
train_data=test_dataset,
|
|
prompt_uris=[
|
|
f"prompts:/{sample_translation_prompt.name}/{sample_translation_prompt.version}"
|
|
],
|
|
scorers=[case_insensitive_match],
|
|
optimizer=testing_optimizer,
|
|
)
|
|
|
|
# Verify custom scorer was used
|
|
# "hola" vs "HOLA" (case insensitive match) -> 1.0
|
|
# "monde" vs "monde" (exact match) -> 1.0
|
|
assert testing_optimizer.captured_scores == [1.0, 1.0]
|
|
assert len(result.optimized_prompts) == 1
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("train_data", "error_match"),
|
|
[
|
|
# Missing inputs validation (handled by _convert_eval_set_to_df)
|
|
([{"outputs": "Hola"}], "Either `inputs` or `trace` column is required"),
|
|
# Empty inputs validation
|
|
(
|
|
[{"inputs": {}, "outputs": "Hola"}],
|
|
"Record 0 is missing required 'inputs' field or it is empty",
|
|
),
|
|
],
|
|
)
|
|
def test_optimize_prompts_validation_errors(
|
|
sample_translation_prompt: PromptVersion,
|
|
train_data: list[dict[str, Any]],
|
|
error_match: str,
|
|
):
|
|
with pytest.raises(MlflowException, match=error_match):
|
|
optimize_prompts(
|
|
predict_fn=sample_predict_fn,
|
|
train_data=train_data,
|
|
prompt_uris=[
|
|
f"prompts:/{sample_translation_prompt.name}/{sample_translation_prompt.version}"
|
|
],
|
|
optimizer=MockPromptOptimizer(),
|
|
scorers=[equivalence],
|
|
)
|
|
|
|
|
|
def test_optimize_prompts_with_chat_prompt(
|
|
sample_translation_prompt: PromptVersion, sample_dataset: pd.DataFrame
|
|
):
|
|
chat_prompt = register_prompt(
|
|
name="test_chat_prompt",
|
|
template=[{"role": "user", "content": "{{input_text}}"}],
|
|
)
|
|
with pytest.raises(MlflowException, match="Only text prompts can be optimized"):
|
|
optimize_prompts(
|
|
predict_fn=sample_predict_fn,
|
|
train_data=sample_dataset,
|
|
prompt_uris=[f"prompts:/{chat_prompt.name}/{chat_prompt.version}"],
|
|
optimizer=MockPromptOptimizer(),
|
|
scorers=[equivalence],
|
|
)
|
|
|
|
|
|
def test_optimize_prompts_with_managed_evaluation_dataset(
|
|
sample_translation_prompt: PromptVersion,
|
|
sample_dataset: pd.DataFrame,
|
|
):
|
|
# Create a `ManagedEvaluationDataset` and populate it with records from sample_dataset
|
|
managed_dataset = create_dataset(name="test_optimize_managed_dataset")
|
|
managed_dataset.merge_records(sample_dataset)
|
|
|
|
result = optimize_prompts(
|
|
predict_fn=sample_predict_fn,
|
|
train_data=managed_dataset,
|
|
prompt_uris=[
|
|
f"prompts:/{sample_translation_prompt.name}/{sample_translation_prompt.version}"
|
|
],
|
|
optimizer=MockPromptOptimizer(),
|
|
scorers=[equivalence],
|
|
)
|
|
|
|
assert len(result.optimized_prompts) == 1
|
|
assert result.initial_eval_score == 0.5
|
|
assert result.final_eval_score == 0.9
|
|
|
|
|
|
def test_optimize_prompts_preserves_model_config(sample_dataset: pd.DataFrame):
|
|
source_model_config = PromptModelConfig(
|
|
provider="openai",
|
|
model_name="gpt-4o",
|
|
temperature=0.7,
|
|
max_tokens=1000,
|
|
)
|
|
|
|
prompt_with_config = register_prompt(
|
|
name="test_prompt_with_model_config",
|
|
template="Translate the following text to {{language}}: {{input_text}}",
|
|
model_config=source_model_config,
|
|
)
|
|
|
|
assert prompt_with_config.model_config is not None
|
|
|
|
def predict_fn(input_text: str, language: str) -> str:
|
|
mlflow.genai.load_prompt(f"prompts:/{prompt_with_config.name}/1")
|
|
translations = {
|
|
("Hello", "Spanish"): "Hola",
|
|
("World", "French"): "Monde",
|
|
("Goodbye", "Spanish"): "Adiós",
|
|
}
|
|
return translations.get((input_text, language), f"translated_{input_text}")
|
|
|
|
result = optimize_prompts(
|
|
predict_fn=predict_fn,
|
|
train_data=sample_dataset,
|
|
prompt_uris=[f"prompts:/{prompt_with_config.name}/{prompt_with_config.version}"],
|
|
optimizer=MockPromptOptimizer(),
|
|
scorers=[equivalence],
|
|
)
|
|
|
|
assert len(result.optimized_prompts) == 1
|
|
optimized_prompt = result.optimized_prompts[0]
|
|
|
|
assert optimized_prompt.model_config["provider"] == "openai"
|
|
assert optimized_prompt.model_config["model_name"] == "gpt-4o"
|
|
assert optimized_prompt.model_config["temperature"] == 0.7
|
|
assert optimized_prompt.model_config["max_tokens"] == 1000
|