stouputils 1.14.0__py3-none-any.whl → 1.14.2__py3-none-any.whl

stouputils/data_science/metric_dictionnary.py

@@ -1,106 +1,106 @@
-"""
-This module contains the MetricDictionnary class, which provides a dictionary of metric names.
-
-This is often used to log metrics to MLflow and to display them in the console easily.
-
-This class contains the following metrics:
-
-1. Main metrics:
-
-   - Area Under the Curve (AUC)
-   - Area Under the Precision-Recall Curve (AUPRC)
-   - Area Under the NPV-Specificity Curve (NEGATIVE_AUPRC)
-   - Specificity (True Negative Rate)
-   - Recall/Sensitivity (True Positive Rate)
-   - Precision (Positive Predictive Value)
-   - Negative Predictive Value (NPV)
-   - Accuracy
-   - F1 Score
-   - Precision-Recall Average
-   - Precision-Recall Average for Negative Class
-
-2. Confusion matrix metrics:
-
-   - True Negatives (TN)
-   - False Positives (FP)
-   - False Negatives (FN)
-   - True Positives (TP)
-   - False Positive Rate
-   - False Negative Rate
-   - False Discovery Rate
-   - False Omission Rate
-   - Critical Success Index (Threat Score)
-
-3. F-scores:
-
-   - F-beta Score (where beta is configurable)
-
-4. Matthews correlation coefficient:
-
-   - Matthews Correlation Coefficient (MCC)
-
-5. Optimal thresholds for binary classification:
-
-   - Youden's J statistic
-   - Cost-based threshold
-   - F1 Score threshold
-   - F1 Score threshold for the negative class
-
-6. Average metrics across folds:
-
-   - Mean value of any metric across k-fold cross validation
-
-7. Standard deviation metrics across folds:
-
-   - Standard deviation of any metric across k-fold cross validation
-"""
-
-class MetricDictionnary:
-
-	# Main metrics (starting with '1:')
-	AUC: str								= "1: Area Under the ROC Curve: AUC / AUROC"
-	AUPRC: str								= "1: Area Under the Precision-Recall Curve: AUPRC / PR AUC"
-	NEGATIVE_AUPRC: str						= "1: Area Under the NPV-Specificity Curve: AUNPRC / NPR AUC"
-	SPECIFICITY: str						= "1: Specificity: True Negative Rate"
-	RECALL: str								= "1: Recall/Sensitivity: True Positive Rate"
-	PRECISION: str							= "1: Precision: Positive Predictive Value"
-	NPV: str								= "1: NPV: Negative Predictive Value"
-	ACCURACY: str							= "1: Accuracy"
-	BALANCED_ACCURACY: str					= "1: Balanced Accuracy"
-	F1_SCORE: str							= "1: F1 Score"
-	F1_SCORE_NEGATIVE: str					= "1: F1 Score for Negative Class"
-	PR_AVERAGE: str							= "1: Precision-Recall Average"
-	PR_AVERAGE_NEGATIVE: str				= "1: Precision-Recall Average for Negative Class"
-
-	# Confusion matrix metrics (starting with '2:')
-	CONFUSION_MATRIX_TN: str				= "2: Confusion Matrix: TN"
-	CONFUSION_MATRIX_FP: str				= "2: Confusion Matrix: FP"
-	CONFUSION_MATRIX_FN: str				= "2: Confusion Matrix: FN"
-	CONFUSION_MATRIX_TP: str				= "2: Confusion Matrix: TP"
-	FALSE_POSITIVE_RATE: str				= "2: False Positive Rate"
-	FALSE_NEGATIVE_RATE: str				= "2: False Negative Rate"
-	FALSE_DISCOVERY_RATE: str				= "2: False Discovery Rate"
-	FALSE_OMISSION_RATE: str				= "2: False Omission Rate"
-	CRITICAL_SUCCESS_INDEX: str				= "2: Critical Success Index: Threat Score"
-
-	# F-scores (starting with '3:')
-	F_SCORE_X: str							= "3: F-X Score"	# X is the beta value
-
-	# Matthews correlation coefficient (starting with '4:')
-	MATTHEWS_CORRELATION_COEFFICIENT: str	= "4: Matthews Correlation Coefficient: MCC"
-
-	# Optimal thresholds (starting with '5:')
-	OPTIMAL_THRESHOLD_YOUDEN: str			= "5: Optimal Threshold: Youden"
-	OPTIMAL_THRESHOLD_COST: str				= "5: Optimal Threshold: Cost"
-	OPTIMAL_THRESHOLD_F1: str				= "5: Optimal Threshold: F1"
-	OPTIMAL_THRESHOLD_F1_NEGATIVE: str		= "5: Optimal Threshold: F1 for Negative Class"
-
-	# Average metrics across folds (starting with '6:')
-	AVERAGE_METRIC: str						= "6: Average METRIC_NAME across folds"
-
-	# Standard deviation metrics across folds (starting with '7:')
-	STANDARD_DEVIATION_METRIC: str			= "7: Standard deviation METRIC_NAME across folds"
-
-	# Parameter finder (starting with '8:')
-	PARAMETER_FINDER: str					= "8: TITLE: PARAMETER_NAME"
-
+"""
+This module contains the MetricDictionnary class, which provides a dictionary of metric names.
+
+This is often used to log metrics to MLflow and to display them in the console easily.
+
+This class contains the following metrics:
+
+1. Main metrics:
+
+   - Area Under the Curve (AUC)
+   - Area Under the Precision-Recall Curve (AUPRC)
+   - Area Under the NPV-Specificity Curve (NEGATIVE_AUPRC)
+   - Specificity (True Negative Rate)
+   - Recall/Sensitivity (True Positive Rate)
+   - Precision (Positive Predictive Value)
+   - Negative Predictive Value (NPV)
+   - Accuracy
+   - F1 Score
+   - Precision-Recall Average
+   - Precision-Recall Average for Negative Class
+
+2. Confusion matrix metrics:
+
+   - True Negatives (TN)
+   - False Positives (FP)
+   - False Negatives (FN)
+   - True Positives (TP)
+   - False Positive Rate
+   - False Negative Rate
+   - False Discovery Rate
+   - False Omission Rate
+   - Critical Success Index (Threat Score)
+
+3. F-scores:
+
+   - F-beta Score (where beta is configurable)
+
+4. Matthews correlation coefficient:
+
+   - Matthews Correlation Coefficient (MCC)
+
+5. Optimal thresholds for binary classification:
+
+   - Youden's J statistic
+   - Cost-based threshold
+   - F1 Score threshold
+   - F1 Score threshold for the negative class
+
+6. Average metrics across folds:
+
+   - Mean value of any metric across k-fold cross validation
+
+7. Standard deviation metrics across folds:
+
+   - Standard deviation of any metric across k-fold cross validation
+"""
+
+class MetricDictionnary:
+
+	# Main metrics (starting with '1:')
+	AUC: str								= "1: Area Under the ROC Curve: AUC / AUROC"
+	AUPRC: str								= "1: Area Under the Precision-Recall Curve: AUPRC / PR AUC"
+	NEGATIVE_AUPRC: str						= "1: Area Under the NPV-Specificity Curve: AUNPRC / NPR AUC"
+	SPECIFICITY: str						= "1: Specificity: True Negative Rate"
+	RECALL: str								= "1: Recall/Sensitivity: True Positive Rate"
+	PRECISION: str							= "1: Precision: Positive Predictive Value"
+	NPV: str								= "1: NPV: Negative Predictive Value"
+	ACCURACY: str							= "1: Accuracy"
+	BALANCED_ACCURACY: str					= "1: Balanced Accuracy"
+	F1_SCORE: str							= "1: F1 Score"
+	F1_SCORE_NEGATIVE: str					= "1: F1 Score for Negative Class"
+	PR_AVERAGE: str							= "1: Precision-Recall Average"
+	PR_AVERAGE_NEGATIVE: str				= "1: Precision-Recall Average for Negative Class"
+
+	# Confusion matrix metrics (starting with '2:')
+	CONFUSION_MATRIX_TN: str				= "2: Confusion Matrix: TN"
+	CONFUSION_MATRIX_FP: str				= "2: Confusion Matrix: FP"
+	CONFUSION_MATRIX_FN: str				= "2: Confusion Matrix: FN"
+	CONFUSION_MATRIX_TP: str				= "2: Confusion Matrix: TP"
+	FALSE_POSITIVE_RATE: str				= "2: False Positive Rate"
+	FALSE_NEGATIVE_RATE: str				= "2: False Negative Rate"
+	FALSE_DISCOVERY_RATE: str				= "2: False Discovery Rate"
+	FALSE_OMISSION_RATE: str				= "2: False Omission Rate"
+	CRITICAL_SUCCESS_INDEX: str				= "2: Critical Success Index: Threat Score"
+
+	# F-scores (starting with '3:')
+	F_SCORE_X: str							= "3: F-X Score"	# X is the beta value
+
+	# Matthews correlation coefficient (starting with '4:')
+	MATTHEWS_CORRELATION_COEFFICIENT: str	= "4: Matthews Correlation Coefficient: MCC"
+
+	# Optimal thresholds (starting with '5:')
+	OPTIMAL_THRESHOLD_YOUDEN: str			= "5: Optimal Threshold: Youden"
+	OPTIMAL_THRESHOLD_COST: str				= "5: Optimal Threshold: Cost"
+	OPTIMAL_THRESHOLD_F1: str				= "5: Optimal Threshold: F1"
+	OPTIMAL_THRESHOLD_F1_NEGATIVE: str		= "5: Optimal Threshold: F1 for Negative Class"
+
+	# Average metrics across folds (starting with '6:')
+	AVERAGE_METRIC: str						= "6: Average METRIC_NAME across folds"
+
+	# Standard deviation metrics across folds (starting with '7:')
+	STANDARD_DEVIATION_METRIC: str			= "7: Standard deviation METRIC_NAME across folds"
+
+	# Parameter finder (starting with '8:')
+	PARAMETER_FINDER: str					= "8: TITLE: PARAMETER_NAME"
+

stouputils/data_science/mlflow_utils.py

@@ -1,206 +1,206 @@
-"""
-This module contains utility functions for working with MLflow.
-
-This module contains functions for:
-
-- Getting the artifact path from the current mlflow run
-- Getting the weights path
-- Getting the runs by experiment name
-- Logging the history of the model to the current mlflow run
-- Starting a new mlflow run
-"""
-
-# Imports
-import os
-from typing import Any, Literal
-
-import mlflow
-from mlflow.entities import Experiment, Run
-
-from ..decorators import handle_error, LogLevels
-from ..io import clean_path
-
-
-# Get artifact path
-def get_artifact_path(from_string: str = "", os_name: str = os.name) -> str:
-	""" Get the artifact path from the current mlflow run (without the file:// prefix).
-
-	Handles the different path formats for Windows and Unix-based systems.
-
-	Args:
-		from_string	(str):		Path to the artifact (optional, defaults to the current mlflow run)
-		os_name		(str):		OS name (optional, defaults to os.name)
-	Returns:
-		str: The artifact path
-	"""
-	# Get the artifact path from the current mlflow run or from a string
-	if not from_string:
-		artifact_path: str = mlflow.get_artifact_uri()
-	else:
-		artifact_path: str = from_string
-
-	# Handle the different path formats for Windows and Unix-based systems
-	if os_name == "nt":
-		return artifact_path.replace("file:///", "")
-	else:
-		return artifact_path.replace("file://", "")
-
-# Get weights path
-def get_weights_path(from_string: str = "", weights_name: str = "best_model.keras", os_name: str = os.name) -> str:
-	""" Get the weights path from the current mlflow run.
-
-	Args:
-		from_string     (str):      Path to the artifact (optional, defaults to the current mlflow run)
-		weights_name    (str):      Name of the weights file (optional, defaults to "best_model.keras")
-		os_name         (str):      OS name (optional, defaults to os.name)
-	Returns:
-		str: The weights path
-
-	Examples:
-		>>> get_weights_path(from_string="file:///path/to/artifact", weights_name="best_model.keras", os_name="posix")
-		'/path/to/artifact/best_model.keras'
-
-		>>> get_weights_path(from_string="file:///C:/path/to/artifact", weights_name="best_model.keras", os_name="nt")
-		'C:/path/to/artifact/best_model.keras'
-	"""
-	return clean_path(f"{get_artifact_path(from_string=from_string, os_name=os_name)}/{weights_name}")
-
-# Get runs by experiment name
-def get_runs_by_experiment_name(experiment_name: str, filter_string: str = "", set_experiment: bool = False) -> list[Run]:
-	""" Get the runs by experiment name.
-
-	Args:
-		experiment_name		(str):		Name of the experiment
-		filter_string		(str):		Filter string to apply to the runs
-		set_experiment		(bool):		Whether to set the experiment
-	Returns:
-		list[Run]:		List of runs
-	"""
-	if set_experiment:
-		mlflow.set_experiment(experiment_name)
-	experiment: Experiment | None = mlflow.get_experiment_by_name(experiment_name)
-	if experiment:
-		return mlflow.search_runs(
-			experiment_ids=[experiment.experiment_id],
-			output_format="list",
-			filter_string=filter_string
-		) # pyright: ignore [reportReturnType]
-	return []
-
-def get_runs_by_model_name(experiment_name: str, model_name: str, set_experiment: bool = False) -> list[Run]:
-	""" Get the runs by model name.
-
-	Args:
-		experiment_name	(str):		Name of the experiment
-		model_name		(str):		Name of the model
-		set_experiment	(bool):		Whether to set the experiment
-	Returns:
-		list[Run]:		List of runs
-	"""
-	return get_runs_by_experiment_name(
-		experiment_name,
-		filter_string=f"tags.model_name = '{model_name}'",
-		set_experiment=set_experiment
-	)
-
-# Log history
-def log_history(history: dict[str, list[Any]], prefix: str = "history", **kwargs: Any) -> None:
-	""" Log the history of the model to the current mlflow run.
-
-	Args:
-		history		(dict[str, list[Any]]):	History of the model
-			(usually from a History object like from a Keras model: history.history)
-		**kwargs	(Any):					Additional arguments to pass to mlflow.log_metric
-	"""
-	for (metric, values) in history.items():
-		for epoch, value in enumerate(values):
-			handle_error(mlflow.log_metric,
-				message=f"Error logging metric {metric}",
-				error_log=LogLevels.ERROR_TRACEBACK
-			)(f"{prefix}_{metric}", value, step=epoch, **kwargs)
-
-
-def start_run(mlflow_uri: str, experiment_name: str, model_name: str, override_run_name: str = "", **kwargs: Any) -> str:
-	""" Start a new mlflow run.
-
-	Args:
-		mlflow_uri			(str):		MLflow URI
-		experiment_name		(str):		Name of the experiment
-		model_name			(str):		Name of the model
-		override_run_name	(str):		Override the run name (if empty, it will be set automatically)
-		**kwargs			(Any):		Additional arguments to pass to mlflow.start_run
-	Returns:
-		str: Name of the run (suffixed with the version number)
-	"""
-	# Set the mlflow URI
-	mlflow.set_tracking_uri(mlflow_uri)
-
-	# Get the runs and increment the version number
-	runs: list[Run] = get_runs_by_model_name(experiment_name, model_name, set_experiment=True)
-	run_number: int = len(runs) + 1
-	run_name: str = f"{model_name}_v{run_number:02d}" if not override_run_name else override_run_name
-
-	# Start the run
-	mlflow.start_run(run_name=run_name, tags={"model_name": model_name}, log_system_metrics=True, **kwargs)
-	return run_name
-
-# Get best run by metric
-def get_best_run_by_metric(
-	experiment_name: str,
-	metric_name: str,
-	model_name: str = "",
-	ascending: bool = False,
-	has_saved_model: bool = True
-) -> Run | None:
-	""" Get the best run by a specific metric.
-
-	Args:
-		experiment_name  (str):    Name of the experiment
-		metric_name      (str):    Name of the metric to sort by
-		model_name       (str):    Name of the model (optional, if empty, all models are considered)
-		ascending        (bool):   Whether to sort in ascending order (default: False, i.e. maximum metric value is best)
-		has_saved_model  (bool):   Whether the model has been saved (default: True)
-	Returns:
-		Run | None: The best run or None if no runs are found
-	"""
-	# Get the runs
-	filter_string: str = f"metrics.`{metric_name}` > 0"
-	if model_name:
-		filter_string += f" AND tags.model_name = '{model_name}'"
-	if has_saved_model:
-		filter_string += " AND tags.has_saved_model = 'True'"
-
-	runs: list[Run] = get_runs_by_experiment_name(
-		experiment_name,
-		filter_string=filter_string,
-		set_experiment=True
-	)
-
-	if not runs:
-		return None
-
-	# Sort the runs by the metric
-	sorted_runs: list[Run] = sorted(
-		runs,
-		key=lambda run: float(run.data.metrics.get(metric_name, 0)), # type: ignore
-		reverse=not ascending
-	)
-
-	return sorted_runs[0] if sorted_runs else None
-
-
-def load_model(run_id: str, model_type: Literal["keras", "pytorch"] = "keras") -> Any:
-	""" Load a model from MLflow.
-
-	Args:
-		run_id      (str):                          ID of the run to load the model from
-		model_type  (Literal["keras", "pytorch"]):  Type of model to load (default: "keras")
-	Returns:
-		Any: The loaded model
-	"""
-	if model_type == "keras":
-		return mlflow.keras.load_model(f"runs:/{run_id}/best_model") # type: ignore
-	elif model_type == "pytorch":
-		return mlflow.pytorch.load_model(f"runs:/{run_id}/best_model") # type: ignore
-	raise ValueError(f"Model type {model_type} not supported")
-
+"""
+This module contains utility functions for working with MLflow.
+
+This module contains functions for:
+
+- Getting the artifact path from the current mlflow run
+- Getting the weights path
+- Getting the runs by experiment name
+- Logging the history of the model to the current mlflow run
+- Starting a new mlflow run
+"""
+
+# Imports
+import os
+from typing import Any, Literal
+
+import mlflow
+from mlflow.entities import Experiment, Run
+
+from ..decorators import handle_error, LogLevels
+from ..io import clean_path
+
+
+# Get artifact path
+def get_artifact_path(from_string: str = "", os_name: str = os.name) -> str:
+	""" Get the artifact path from the current mlflow run (without the file:// prefix).
+
+	Handles the different path formats for Windows and Unix-based systems.
+
+	Args:
+		from_string	(str):		Path to the artifact (optional, defaults to the current mlflow run)
+		os_name		(str):		OS name (optional, defaults to os.name)
+	Returns:
+		str: The artifact path
+	"""
+	# Get the artifact path from the current mlflow run or from a string
+	if not from_string:
+		artifact_path: str = mlflow.get_artifact_uri()
+	else:
+		artifact_path: str = from_string
+
+	# Handle the different path formats for Windows and Unix-based systems
+	if os_name == "nt":
+		return artifact_path.replace("file:///", "")
+	else:
+		return artifact_path.replace("file://", "")
+
+# Get weights path
+def get_weights_path(from_string: str = "", weights_name: str = "best_model.keras", os_name: str = os.name) -> str:
+	""" Get the weights path from the current mlflow run.
+
+	Args:
+		from_string     (str):      Path to the artifact (optional, defaults to the current mlflow run)
+		weights_name    (str):      Name of the weights file (optional, defaults to "best_model.keras")
+		os_name         (str):      OS name (optional, defaults to os.name)
+	Returns:
+		str: The weights path
+
+	Examples:
+		>>> get_weights_path(from_string="file:///path/to/artifact", weights_name="best_model.keras", os_name="posix")
+		'/path/to/artifact/best_model.keras'
+
+		>>> get_weights_path(from_string="file:///C:/path/to/artifact", weights_name="best_model.keras", os_name="nt")
+		'C:/path/to/artifact/best_model.keras'
+	"""
+	return clean_path(f"{get_artifact_path(from_string=from_string, os_name=os_name)}/{weights_name}")
+
+# Get runs by experiment name
+def get_runs_by_experiment_name(experiment_name: str, filter_string: str = "", set_experiment: bool = False) -> list[Run]:
+	""" Get the runs by experiment name.
+
+	Args:
+		experiment_name		(str):		Name of the experiment
+		filter_string		(str):		Filter string to apply to the runs
+		set_experiment		(bool):		Whether to set the experiment
+	Returns:
+		list[Run]:		List of runs
+	"""
+	if set_experiment:
+		mlflow.set_experiment(experiment_name)
+	experiment: Experiment | None = mlflow.get_experiment_by_name(experiment_name)
+	if experiment:
+		return mlflow.search_runs(
+			experiment_ids=[experiment.experiment_id],
+			output_format="list",
+			filter_string=filter_string
+		) # pyright: ignore [reportReturnType]
+	return []
+
+def get_runs_by_model_name(experiment_name: str, model_name: str, set_experiment: bool = False) -> list[Run]:
+	""" Get the runs by model name.
+
+	Args:
+		experiment_name	(str):		Name of the experiment
+		model_name		(str):		Name of the model
+		set_experiment	(bool):		Whether to set the experiment
+	Returns:
+		list[Run]:		List of runs
+	"""
+	return get_runs_by_experiment_name(
+		experiment_name,
+		filter_string=f"tags.model_name = '{model_name}'",
+		set_experiment=set_experiment
+	)
+
+# Log history
+def log_history(history: dict[str, list[Any]], prefix: str = "history", **kwargs: Any) -> None:
+	""" Log the history of the model to the current mlflow run.
+
+	Args:
+		history		(dict[str, list[Any]]):	History of the model
+			(usually from a History object like from a Keras model: history.history)
+		**kwargs	(Any):					Additional arguments to pass to mlflow.log_metric
+	"""
+	for (metric, values) in history.items():
+		for epoch, value in enumerate(values):
+			handle_error(mlflow.log_metric,
+				message=f"Error logging metric {metric}",
+				error_log=LogLevels.ERROR_TRACEBACK
+			)(f"{prefix}_{metric}", value, step=epoch, **kwargs)
+
+
+def start_run(mlflow_uri: str, experiment_name: str, model_name: str, override_run_name: str = "", **kwargs: Any) -> str:
+	""" Start a new mlflow run.
+
+	Args:
+		mlflow_uri			(str):		MLflow URI
+		experiment_name		(str):		Name of the experiment
+		model_name			(str):		Name of the model
+		override_run_name	(str):		Override the run name (if empty, it will be set automatically)
+		**kwargs			(Any):		Additional arguments to pass to mlflow.start_run
+	Returns:
+		str: Name of the run (suffixed with the version number)
+	"""
+	# Set the mlflow URI
+	mlflow.set_tracking_uri(mlflow_uri)
+
+	# Get the runs and increment the version number
+	runs: list[Run] = get_runs_by_model_name(experiment_name, model_name, set_experiment=True)
+	run_number: int = len(runs) + 1
+	run_name: str = f"{model_name}_v{run_number:02d}" if not override_run_name else override_run_name
+
+	# Start the run
+	mlflow.start_run(run_name=run_name, tags={"model_name": model_name}, log_system_metrics=True, **kwargs)
+	return run_name
+
+# Get best run by metric
+def get_best_run_by_metric(
+	experiment_name: str,
+	metric_name: str,
+	model_name: str = "",
+	ascending: bool = False,
+	has_saved_model: bool = True
+) -> Run | None:
+	""" Get the best run by a specific metric.
+
+	Args:
+		experiment_name  (str):    Name of the experiment
+		metric_name      (str):    Name of the metric to sort by
+		model_name       (str):    Name of the model (optional, if empty, all models are considered)
+		ascending        (bool):   Whether to sort in ascending order (default: False, i.e. maximum metric value is best)
+		has_saved_model  (bool):   Whether the model has been saved (default: True)
+	Returns:
+		Run | None: The best run or None if no runs are found
+	"""
+	# Get the runs
+	filter_string: str = f"metrics.`{metric_name}` > 0"
+	if model_name:
+		filter_string += f" AND tags.model_name = '{model_name}'"
+	if has_saved_model:
+		filter_string += " AND tags.has_saved_model = 'True'"
+
+	runs: list[Run] = get_runs_by_experiment_name(
+		experiment_name,
+		filter_string=filter_string,
+		set_experiment=True
+	)
+
+	if not runs:
+		return None
+
+	# Sort the runs by the metric
+	sorted_runs: list[Run] = sorted(
+		runs,
+		key=lambda run: float(run.data.metrics.get(metric_name, 0)), # type: ignore
+		reverse=not ascending
+	)
+
+	return sorted_runs[0] if sorted_runs else None
+
+
+def load_model(run_id: str, model_type: Literal["keras", "pytorch"] = "keras") -> Any:
+	""" Load a model from MLflow.
+
+	Args:
+		run_id      (str):                          ID of the run to load the model from
+		model_type  (Literal["keras", "pytorch"]):  Type of model to load (default: "keras")
+	Returns:
+		Any: The loaded model
+	"""
+	if model_type == "keras":
+		return mlflow.keras.load_model(f"runs:/{run_id}/best_model") # type: ignore
+	elif model_type == "pytorch":
+		return mlflow.pytorch.load_model(f"runs:/{run_id}/best_model") # type: ignore
+	raise ValueError(f"Model type {model_type} not supported")
+