Spaces:
Sleeping
Sleeping
File size: 2,247 Bytes
72e2b6e | 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 | from pathlib import Path
import matplotlib.pyplot as plt
import mlflow
import numpy as np
from sklearn.metrics import ConfusionMatrixDisplay, confusion_matrix
def get_or_create_experiment(name: str) -> str:
experiment = mlflow.get_experiment_by_name(name)
if experiment is None:
return mlflow.create_experiment(name)
return experiment.experiment_id
def log_config(config: dict) -> None:
flat = {}
for section, values in config.items():
if isinstance(values, dict):
for k, v in values.items():
flat[f"{section}.{k}"] = v
else:
flat[section] = values
mlflow.log_params(flat)
def log_metrics(metrics: dict, step: int | None = None) -> None:
mlflow.log_metrics(metrics, step=step)
def log_confusion_matrix(
y_true: np.ndarray,
y_pred: np.ndarray,
labels: list[str],
save_path: str = "artifacts/confusion_matrix.png",
) -> None:
path = Path(save_path)
path.parent.mkdir(parents=True, exist_ok=True)
cm = confusion_matrix(y_true, y_pred)
fig, ax = plt.subplots(figsize=(20, 20))
disp = ConfusionMatrixDisplay(confusion_matrix=cm, display_labels=labels)
disp.plot(ax=ax, xticks_rotation=90, colorbar=False)
ax.set_title("Confusion Matrix")
plt.tight_layout()
plt.savefig(path, dpi=100)
plt.close()
mlflow.log_artifact(str(path))
def register_model(run_id: str, artifact_path: str, model_name: str, alias: str = "champion") -> None:
model_uri = f"runs:/{run_id}/{artifact_path}"
result = mlflow.register_model(model_uri, model_name)
client = mlflow.MlflowClient()
client.set_registered_model_alias(
name=model_name,
alias=alias,
version=result.version,
)
print(f" registered {model_name} v{result.version} with alias '{alias}'")
def load_registered_model(model_name: str, alias: str = "champion"):
model_uri = f"models:/{model_name}@{alias}"
return mlflow.pyfunc.load_model(model_uri)
def get_model_version_by_alias(model_name: str, alias: str = "champion"):
client = mlflow.MlflowClient()
return client.get_model_version_by_alias(model_name, alias)
|