stouputils 1.14.0__py3-none-any.whl

stouputils/data_science/mlflow_utils.py

@@ -0,0 +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")
+

stouputils/data_science/models/abstract_model.py

@@ -0,0 +1,149 @@
+""" Abstract base class for all model implementations.
+Defines the interface that all concrete model classes must implement.
+
+Provides abstract methods for core model operations including:
+
+- Class routine management
+- Model loading
+- Training procedures
+- Prediction functionality
+- Evaluation metrics
+
+Classes inheriting from AbstractModel must implement all methods.
+"""
+
+# Imports
+from __future__ import annotations
+
+import multiprocessing.queues
+from collections.abc import Iterable
+from tempfile import TemporaryDirectory
+from typing import Any
+
+from ...decorators import abstract, LogLevels
+
+from ..dataset import Dataset
+
+
+# Base class
+class AbstractModel:
+	""" Abstract class for all models to copy and implement the methods. """
+	# Class methods
+	@abstract(error_log=LogLevels.ERROR_TRACEBACK)
+	def __init__(
+		self, num_classes: int, kfold: int = 0, transfer_learning: str = "imagenet", **override_params: Any
+	) -> None:
+		pass
+
+
+	## Public abstract methods
+	@abstract(error_log=LogLevels.ERROR_TRACEBACK)
+	def routine_full(self, dataset: Dataset, verbose: int = 0) -> AbstractModel:
+		return self
+
+	@abstract(error_log=LogLevels.ERROR_TRACEBACK)
+	def class_load(self) -> None:
+		pass
+
+
+	@abstract(error_log=LogLevels.ERROR_TRACEBACK)
+	def class_train(self, dataset: Dataset, verbose: int = 0) -> bool:
+		return False
+
+
+	@abstract(error_log=LogLevels.ERROR_TRACEBACK)
+	def class_predict(self, X_test: Iterable[Any]) -> Iterable[Any]:
+		return []
+
+
+	@abstract(error_log=LogLevels.ERROR_TRACEBACK)
+	def class_evaluate(
+		self,
+		dataset: Dataset,
+		metrics_names: tuple[str, ...] = (),
+		save_model: bool = False,
+		verbose: int = 0
+	) -> bool:
+		return False
+
+
+	## Protected abstract methods
+	@abstract(error_log=LogLevels.ERROR_TRACEBACK)
+	def _fit(
+		self,
+		model: Any,
+		x: Any,
+		y: Any | None = None,
+		validation_data: tuple[Any, Any] | None = None,
+		shuffle: bool = True,
+		batch_size: int | None = None,
+		epochs: int = 1,
+		callbacks: list[Any] | None = None,
+		class_weight: dict[int, float] | None = None,
+		verbose: int = 0,
+		*args: Any,
+		**kwargs: Any
+	) -> Any:
+		pass
+
+	@abstract(error_log=LogLevels.ERROR_TRACEBACK)
+	def _get_callbacks(self) -> list[Any]:
+		return []
+
+	@abstract(error_log=LogLevels.ERROR_TRACEBACK)
+	def _get_metrics(self) -> list[Any]:
+		return []
+
+	@abstract(error_log=LogLevels.ERROR_TRACEBACK)
+	def _get_optimizer(self, learning_rate: float = 0.0) -> Any:
+		pass
+
+	@abstract(error_log=LogLevels.ERROR_TRACEBACK)
+	def _get_loss(self) -> Any:
+		pass
+
+	@abstract(error_log=LogLevels.ERROR_TRACEBACK)
+	def _get_base_model(self) -> Any:
+		pass
+
+	@abstract(error_log=LogLevels.ERROR_TRACEBACK)
+	def _get_architectures(
+		self, optimizer: Any = None, loss: Any = None, metrics: list[Any] | None = None
+	) -> tuple[Any, Any]:
+		return (None, None)
+
+	@abstract(error_log=LogLevels.ERROR_TRACEBACK)
+	def _find_best_learning_rate(self, dataset: Dataset, verbose: int = 0) -> float:
+		return 0.0
+
+	@abstract(error_log=LogLevels.ERROR_TRACEBACK)
+	def _train_fold(self, dataset: Dataset, fold_number: int = 0, mlflow_prefix: str = "history", verbose: int = 0) -> Any:
+		pass
+
+	@abstract(error_log=LogLevels.ERROR_TRACEBACK)
+	def _log_final_model(self) -> None:
+		pass
+
+	@abstract(error_log=LogLevels.ERROR_TRACEBACK)
+	def _find_best_learning_rate_subprocess(
+		self, dataset: Dataset, queue: multiprocessing.queues.Queue[Any] | None = None, verbose: int = 0
+	) -> dict[str, Any] | None:
+		pass
+
+	@abstract(error_log=LogLevels.ERROR_TRACEBACK)
+	def _find_best_unfreeze_percentage_subprocess(
+		self, dataset: Dataset, queue: multiprocessing.queues.Queue[Any] | None = None, verbose: int = 0
+	) -> dict[str, Any] | None:
+		pass
+
+	@abstract(error_log=LogLevels.ERROR_TRACEBACK)
+	def _train_subprocess(
+		self,
+		dataset: Dataset,
+		checkpoint_path: str,
+		temp_dir: TemporaryDirectory[str] | None = None,
+		queue: multiprocessing.queues.Queue[Any] | None = None,
+		verbose: int = 0
+	) -> dict[str, Any] | None:
+		pass
+

stouputils/data_science/models/all.py

@@ -0,0 +1,85 @@
+
+# Imports
+import itertools
+
+from .keras.all import (
+    VGG16,
+    VGG19,
+    ConvNeXtBase,
+    ConvNeXtLarge,
+    ConvNeXtSmall,
+    ConvNeXtTiny,
+    ConvNeXtXLarge,
+    DenseNet121,
+    DenseNet169,
+    DenseNet201,
+    EfficientNetB0,
+    EfficientNetV2B0,
+    EfficientNetV2L,
+    EfficientNetV2M,
+    EfficientNetV2S,
+    MobileNet,
+    MobileNetV2,
+    MobileNetV3Large,
+    MobileNetV3Small,
+    ResNet50V2,
+    ResNet101V2,
+    ResNet152V2,
+    SqueezeNet,
+    Xception,
+)
+from .model_interface import ModelInterface
+
+# Other models
+from .sandbox import Sandbox
+
+
+# Create a custom dictionary class to allow for documentation
+class ModelClassMap(dict[type[ModelInterface], tuple[str, ...]]):
+    pass
+
+# Routine map
+CLASS_MAP: ModelClassMap = ModelClassMap({
+    SqueezeNet:         ("squeezenet", "squeezenets", "all", "often"),
+
+    DenseNet121:        ("densenet121", "densenets", "all", "often", "good"),
+    DenseNet169:        ("densenet169", "densenets", "all", "often", "good"),
+    DenseNet201:        ("densenet201", "densenets", "all", "often", "good"),
+
+    EfficientNetB0:     ("efficientnetb0", "efficientnets", "all"),
+    EfficientNetV2B0:   ("efficientnetv2b0", "efficientnets", "all"),
+    EfficientNetV2S:    ("efficientnetv2s", "efficientnets", "all", "often"),
+    EfficientNetV2M:    ("efficientnetv2m", "efficientnets", "all", "often"),
+    EfficientNetV2L:    ("efficientnetv2l", "efficientnets", "all", "often"),
+
+    ConvNeXtTiny:       ("convnexttiny", "convnexts", "all", "often", "good"),
+    ConvNeXtSmall:      ("convnextsmall", "convnexts", "all", "often"),
+    ConvNeXtBase:       ("convnextbase", "convnexts", "all", "often", "good"),
+    ConvNeXtLarge:      ("convnextlarge", "convnexts", "all", "often"),
+    ConvNeXtXLarge:     ("convnextxlarge", "convnexts", "all", "often", "good"),
+
+    VGG16:              ("vgg16", "vggs", "all"),
+    VGG19:              ("vgg19", "vggs", "all"),
+
+    MobileNet:          ("mobilenet", "mobilenets", "all"),
+    MobileNetV2:        ("mobilenetv2", "mobilenets", "all", "often"),
+    MobileNetV3Small:   ("mobilenetv3small", "mobilenets", "all", "often"),
+    MobileNetV3Large:   ("mobilenetv3large", "mobilenets", "all", "often", "good"),
+
+    ResNet50V2:         ("resnet50v2", "resnetsv2", "resnets", "all", "often"),
+    ResNet101V2:        ("resnet101v2", "resnetsv2", "resnets", "all", "often"),
+    ResNet152V2:        ("resnet152v2", "resnetsv2", "resnets", "all", "often"),
+
+    Xception:           ("xception", "xceptions", "all", "often"),
+    Sandbox:            ("sandbox",),
+})
+
+# All models names and aliases
+ALL_MODELS: list[str] = sorted(set(itertools.chain.from_iterable(v for v in CLASS_MAP.values())))
+""" All models names and aliases found in the `CLASS_MAP` dictionary. """
+
+# Additional docstring
+new_docstring: str = "\n\n" + "\n".join(f"- {k.__name__}: {v}" for k, v in CLASS_MAP.items())
+ModelClassMap.__doc__ = "Dictionary mapping class to their names and aliases. " + new_docstring
+CLASS_MAP.__doc__ = ModelClassMap.__doc__
+