File size: 5,118 Bytes
243b4bc | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 | import os
import pandas as pd
import pytest
import torch
from unittest.mock import MagicMock, patch
from src.components.model_evaluation import Model_Evaluation
from src.entity.config_entity import ModelEvaluationConfig, ModelTrainingConfig
from src.entity.artifact_entity import DataTransformationArtifact, ModelTrainingArtifact, ModelEvaluationArtifact
@pytest.mark.asyncio
async def test_model_evaluation_initiate_success(dummy_dataframe_transformed, temp_artifact_dir):
"""Test that Model_Evaluation successfully runs evaluation on test data,
saves metrics, plots, and logs to MLflow (mocked) without hitting remote servers.
"""
# 1. Prepare file paths in the temporary artifact directory
train_path = os.path.join(temp_artifact_dir, "train.csv")
val_path = os.path.join(temp_artifact_dir, "val.csv")
test_path = os.path.join(temp_artifact_dir, "test.csv")
images_path = os.path.join(temp_artifact_dir, "images")
# Save the transformed dataframe to test_path
dummy_dataframe_transformed.to_csv(test_path, index=False)
# 2. Create mock artifacts
data_transformation_artifact = DataTransformationArtifact(
train_path=train_path,
test_path=test_path,
val_path=val_path,
images_path=images_path
)
model_training_artifact = ModelTrainingArtifact(
model_path=os.path.join(temp_artifact_dir, "checkpoints", "test_model.pt"),
is_trained=True,
message="Model trained"
)
# 3. Setup configurations
eval_dir = os.path.join(temp_artifact_dir, "model_evaluation")
model_evaluation_config = ModelEvaluationConfig(
evaluation_artifact_dir=eval_dir,
metrics_file_name="metrics.yaml",
loss_plot_file_name="loss_plot.png",
confusion_matrix_file_name="confusion_matrix.png",
mlflow_experiment_name="test-experiment",
mlflow_run_name="test-run"
)
cache_dir = os.path.join(temp_artifact_dir, "cache")
os.makedirs(cache_dir, exist_ok=True)
model_training_config = ModelTrainingConfig(
epochs=1,
patience=1,
batch_size=2,
model_name="test_model.pt",
model_dir=os.path.join(temp_artifact_dir, "checkpoints"),
cache_dir=cache_dir
)
# 4. Pre-create the cached features to skip actual image/text encoders
num_samples = len(dummy_dataframe_transformed)
img_feats = torch.randn((num_samples, model_training_config.image_feature_output))
txt_feats = torch.randn((num_samples, model_training_config.text_feature_output))
labels = torch.tensor(dummy_dataframe_transformed['label'].values, dtype=torch.float32)
torch.save(img_feats, os.path.join(cache_dir, "img_feats.pt"))
torch.save(txt_feats, os.path.join(cache_dir, "txt_feats.pt"))
torch.save(labels, os.path.join(cache_dir, "labels.pt"))
# 5. Patch dagshub and mlflow to prevent any remote API requests/hosting
with patch("dagshub.auth.add_app_token") as mock_add_token, \
patch("dagshub.init") as mock_dag_init, \
patch("src.components.model_evaluation.mlflow") as mock_mlflow:
# Setup mlflow mock context manager for start_run
mock_run = MagicMock()
mock_mlflow.start_run.return_value = mock_run
# Patch load_model to bypass reading checkpoint from disk
with patch("src.entity.model.MyModel.load_model") as mock_load_model:
evaluator = Model_Evaluation(
model_evaluation_config=model_evaluation_config,
data_transformation_artifact=data_transformation_artifact,
model_training_artifact=model_training_artifact,
model_training_config=model_training_config
)
# Set dummy losses on the model to test loss plotting and log extraction
evaluator.model.train_loss = [0.5, 0.4]
evaluator.model.val_loss = [0.6, 0.45]
# 6. Run evaluation
artifact = await evaluator.initiate()
# 7. Assertions on output artifact and generated files
assert isinstance(artifact, ModelEvaluationArtifact)
assert artifact.is_evaluated is True
assert os.path.isfile(artifact.metrics_file_path)
assert os.path.isfile(artifact.loss_plot_path)
assert os.path.isfile(artifact.confusion_matrix_path)
# Verify MLflow interaction was triggered locally via mock
mock_mlflow.set_experiment.assert_called_once()
mock_mlflow.start_run.assert_called_once_with(run_name=model_evaluation_config.mlflow_run_name)
mock_mlflow.log_params.assert_called_once()
# Verify yaml metrics structure
with open(artifact.metrics_file_path, "r") as f:
import yaml
metrics = yaml.safe_load(f)
for key in ["train_loss", "val_loss", "test_accuracy", "test_precision", "test_recall", "test_f1_score"]:
assert key in metrics
|