mlrun 1.6.4rc8__py3-none-any.whl → 1.7.0__py3-none-any.whl

mlrun/frameworks/pytorch/callbacks_handler.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
-from typing import Dict, List, Tuple, Union
+from typing import Union
 
 from torch import Tensor
 from torch.nn import Module
@@ -66,7 +66,7 @@ class CallbacksHandler:
     A class for handling multiple callbacks during a run.
     """
 
-    def __init__(self, callbacks: List[Union[Callback, Tuple[str, Callback]]]):
+    def __init__(self, callbacks: list[Union[Callback, tuple[str, Callback]]]):
         """
         Initialize the callbacks handler with the given callbacks he will handle. The callbacks can be passed as their
         initialized instances or as a tuple where [0] is a name that will be attached to him and [1] will be the
@@ -99,7 +99,7 @@ class CallbacksHandler:
                 self._callbacks[callback.__class__.__name__] = callback
 
     @property
-    def callbacks(self) -> Dict[str, Callback]:
+    def callbacks(self) -> dict[str, Callback]:
         """
         Get the callbacks dictionary handled by this handler.
 
@@ -114,9 +114,9 @@ class CallbacksHandler:
         validation_set: DataLoader,
         loss_function: Module,
         optimizer: Optimizer,
-        metric_functions: List[PyTorchTypes.MetricFunctionType],
+        metric_functions: list[PyTorchTypes.MetricFunctionType],
         scheduler,
-        callbacks: List[str] = None,
+        callbacks: list[str] = None,
     ) -> bool:
         """
         Call the 'on_setup' method of every callback in the callbacks list. If the list is 'None' (not given), all
@@ -145,7 +145,7 @@ class CallbacksHandler:
             scheduler=scheduler,
         )
 
-    def on_run_begin(self, callbacks: List[str] = None) -> bool:
+    def on_run_begin(self, callbacks: list[str] = None) -> bool:
         """
         Call the 'on_run_begin' method of every callback in the callbacks list. If the list is 'None' (not given), all
         callbacks will be called.
@@ -159,7 +159,7 @@ class CallbacksHandler:
             callbacks=self._parse_names(names=callbacks),
         )
 
-    def on_run_end(self, callbacks: List[str] = None) -> bool:
+    def on_run_end(self, callbacks: list[str] = None) -> bool:
         """
         Call the 'on_run_end' method of every callback in the callbacks list. If the list is 'None' (not given), all
         callbacks will be called.
@@ -173,7 +173,7 @@ class CallbacksHandler:
             callbacks=self._parse_names(names=callbacks),
         )
 
-    def on_epoch_begin(self, epoch: int, callbacks: List[str] = None) -> bool:
+    def on_epoch_begin(self, epoch: int, callbacks: list[str] = None) -> bool:
         """
         Call the 'on_epoch_begin' method of every callback in the callbacks list. If the list is 'None' (not given), all
         callbacks will be called.
@@ -189,7 +189,7 @@ class CallbacksHandler:
             epoch=epoch,
         )
 
-    def on_epoch_end(self, epoch: int, callbacks: List[str] = None) -> bool:
+    def on_epoch_end(self, epoch: int, callbacks: list[str] = None) -> bool:
         """
         Call the 'on_epoch_end' method of every callback in the callbacks list. If the list is 'None' (not given), all
         callbacks will be called.
@@ -205,7 +205,7 @@ class CallbacksHandler:
             epoch=epoch,
         )
 
-    def on_train_begin(self, callbacks: List[str] = None) -> bool:
+    def on_train_begin(self, callbacks: list[str] = None) -> bool:
         """
         Call the 'on_train_begin' method of every callback in the callbacks list. If the list is 'None' (not given), all
         callbacks will be called.
@@ -219,7 +219,7 @@ class CallbacksHandler:
             callbacks=self._parse_names(names=callbacks),
         )
 
-    def on_train_end(self, callbacks: List[str] = None) -> bool:
+    def on_train_end(self, callbacks: list[str] = None) -> bool:
         """
         Call the 'on_train_end' method of every callback in the callbacks list. If the list is 'None' (not given), all
         callbacks will be called.
@@ -233,7 +233,7 @@ class CallbacksHandler:
             callbacks=self._parse_names(names=callbacks),
         )
 
-    def on_validation_begin(self, callbacks: List[str] = None) -> bool:
+    def on_validation_begin(self, callbacks: list[str] = None) -> bool:
         """
         Call the 'on_validation_begin' method of every callback in the callbacks list. If the list is 'None'
         (not given), all callbacks will be called.
@@ -250,8 +250,8 @@ class CallbacksHandler:
     def on_validation_end(
         self,
         loss_value: PyTorchTypes.MetricValueType,
-        metric_values: List[float],
-        callbacks: List[str] = None,
+        metric_values: list[float],
+        callbacks: list[str] = None,
     ) -> bool:
         """
         Call the 'on_validation_end' method of every callback in the callbacks list. If the list is 'None' (not given),
@@ -271,7 +271,7 @@ class CallbacksHandler:
         )
 
     def on_train_batch_begin(
-        self, batch: int, x, y_true: Tensor, callbacks: List[str] = None
+        self, batch: int, x, y_true: Tensor, callbacks: list[str] = None
     ) -> bool:
         """
         Call the 'on_train_batch_begin' method of every callback in the callbacks list. If the list is 'None'
@@ -298,7 +298,7 @@ class CallbacksHandler:
         x,
         y_pred: Tensor,
         y_true: Tensor,
-        callbacks: List[str] = None,
+        callbacks: list[str] = None,
     ) -> bool:
         """
         Call the 'on_train_batch_end' method of every callback in the callbacks list. If the list is 'None' (not given),
@@ -322,7 +322,7 @@ class CallbacksHandler:
         )
 
     def on_validation_batch_begin(
-        self, batch: int, x, y_true: Tensor, callbacks: List[str] = None
+        self, batch: int, x, y_true: Tensor, callbacks: list[str] = None
     ) -> bool:
         """
         Call the 'on_validation_batch_begin' method of every callback in the callbacks list. If the list is 'None'
@@ -349,7 +349,7 @@ class CallbacksHandler:
         x,
         y_pred: Tensor,
         y_true: Tensor,
-        callbacks: List[str] = None,
+        callbacks: list[str] = None,
     ) -> bool:
         """
         Call the 'on_validation_batch_end' method of every callback in the callbacks list. If the list is 'None'
@@ -375,7 +375,7 @@ class CallbacksHandler:
     def on_inference_begin(
         self,
         x,
-        callbacks: List[str] = None,
+        callbacks: list[str] = None,
     ) -> bool:
         """
         Call the 'on_inference_begin' method of every callback in the callbacks list. If the list is 'None' (not given),
@@ -396,7 +396,7 @@ class CallbacksHandler:
         self,
         y_pred: Tensor,
         y_true: Tensor,
-        callbacks: List[str] = None,
+        callbacks: list[str] = None,
     ) -> bool:
         """
         Call the 'on_inference_end' method of every callback in the callbacks list. If the list is 'None' (not given),
@@ -415,7 +415,7 @@ class CallbacksHandler:
             y_true=y_true,
         )
 
-    def on_train_loss_begin(self, callbacks: List[str] = None) -> bool:
+    def on_train_loss_begin(self, callbacks: list[str] = None) -> bool:
         """
         Call the 'on_train_loss_begin' method of every callback in the callbacks list. If the list is 'None'
         (not given), all callbacks will be called.
@@ -430,7 +430,7 @@ class CallbacksHandler:
         )
 
     def on_train_loss_end(
-        self, loss_value: PyTorchTypes.MetricValueType, callbacks: List[str] = None
+        self, loss_value: PyTorchTypes.MetricValueType, callbacks: list[str] = None
     ) -> bool:
         """
         Call the 'on_train_loss_end' method of every callback in the callbacks list. If the list is 'None' (not given),
@@ -447,7 +447,7 @@ class CallbacksHandler:
             loss_value=loss_value,
         )
 
-    def on_validation_loss_begin(self, callbacks: List[str] = None) -> bool:
+    def on_validation_loss_begin(self, callbacks: list[str] = None) -> bool:
         """
         Call the 'on_validation_loss_begin' method of every callback in the callbacks list. If the list is 'None'
         (not given), all callbacks will be called.
@@ -462,7 +462,7 @@ class CallbacksHandler:
         )
 
     def on_validation_loss_end(
-        self, loss_value: PyTorchTypes.MetricValueType, callbacks: List[str] = None
+        self, loss_value: PyTorchTypes.MetricValueType, callbacks: list[str] = None
     ) -> bool:
         """
         Call the 'on_validation_loss_end' method of every callback in the callbacks list. If the list is 'None'
@@ -479,7 +479,7 @@ class CallbacksHandler:
             loss_value=loss_value,
         )
 
-    def on_train_metrics_begin(self, callbacks: List[str] = None) -> bool:
+    def on_train_metrics_begin(self, callbacks: list[str] = None) -> bool:
         """
         Call the 'on_train_metrics_begin' method of every callback in the callbacks list. If the list is 'None'
         (not given), all callbacks will be called.
@@ -495,8 +495,8 @@ class CallbacksHandler:
 
     def on_train_metrics_end(
         self,
-        metric_values: List[PyTorchTypes.MetricValueType],
-        callbacks: List[str] = None,
+        metric_values: list[PyTorchTypes.MetricValueType],
+        callbacks: list[str] = None,
     ) -> bool:
         """
         Call the 'on_train_metrics_end' method of every callback in the callbacks list. If the list is 'None'
@@ -513,7 +513,7 @@ class CallbacksHandler:
             metric_values=metric_values,
         )
 
-    def on_validation_metrics_begin(self, callbacks: List[str] = None) -> bool:
+    def on_validation_metrics_begin(self, callbacks: list[str] = None) -> bool:
         """
         Call the 'on_validation_metrics_begin' method of every callback in the callbacks list. If the list is 'None'
         (not given), all callbacks will be called.
@@ -529,8 +529,8 @@ class CallbacksHandler:
 
     def on_validation_metrics_end(
         self,
-        metric_values: List[PyTorchTypes.MetricValueType],
-        callbacks: List[str] = None,
+        metric_values: list[PyTorchTypes.MetricValueType],
+        callbacks: list[str] = None,
     ) -> bool:
         """
         Call the 'on_validation_metrics_end' method of every callback in the callbacks list. If the list is 'None'
@@ -547,7 +547,7 @@ class CallbacksHandler:
             metric_values=metric_values,
         )
 
-    def on_backward_begin(self, callbacks: List[str] = None) -> bool:
+    def on_backward_begin(self, callbacks: list[str] = None) -> bool:
         """
         Call the 'on_backward_begin' method of every callback in the callbacks list. If the list is 'None' (not given),
         all callbacks will be called.
@@ -561,7 +561,7 @@ class CallbacksHandler:
             callbacks=self._parse_names(names=callbacks),
         )
 
-    def on_backward_end(self, callbacks: List[str] = None) -> bool:
+    def on_backward_end(self, callbacks: list[str] = None) -> bool:
         """
         Call the 'on_backward_end' method of every callback in the callbacks list. If the list is 'None' (not given),
         all callbacks will be called.
@@ -575,7 +575,7 @@ class CallbacksHandler:
             callbacks=self._parse_names(names=callbacks),
         )
 
-    def on_optimizer_step_begin(self, callbacks: List[str] = None) -> bool:
+    def on_optimizer_step_begin(self, callbacks: list[str] = None) -> bool:
         """
         Call the 'on_optimizer_step_begin' method of every callback in the callbacks list. If the list is 'None'
         (not given), all callbacks will be called.
@@ -589,7 +589,7 @@ class CallbacksHandler:
             callbacks=self._parse_names(names=callbacks),
         )
 
-    def on_optimizer_step_end(self, callbacks: List[str] = None) -> bool:
+    def on_optimizer_step_end(self, callbacks: list[str] = None) -> bool:
         """
         Call the 'on_optimizer_step_end' method of every callback in the callbacks list. If the list is 'None'
         (not given), all callbacks will be called.
@@ -603,7 +603,7 @@ class CallbacksHandler:
             callbacks=self._parse_names(names=callbacks),
         )
 
-    def on_scheduler_step_begin(self, callbacks: List[str] = None) -> bool:
+    def on_scheduler_step_begin(self, callbacks: list[str] = None) -> bool:
         """
         Call the 'on_scheduler_step_begin' method of every callback in the callbacks list. If the list is 'None'
         (not given), all callbacks will be called.
@@ -617,7 +617,7 @@ class CallbacksHandler:
             callbacks=self._parse_names(names=callbacks),
         )
 
-    def on_scheduler_step_end(self, callbacks: List[str] = None) -> bool:
+    def on_scheduler_step_end(self, callbacks: list[str] = None) -> bool:
         """
         Call the 'on_scheduler_step_end' method of every callback in the callbacks list. If the list is 'None'
         (not given), all callbacks will be called.
@@ -631,7 +631,7 @@ class CallbacksHandler:
             callbacks=self._parse_names(names=callbacks),
         )
 
-    def _parse_names(self, names: Union[List[str], None]) -> List[str]:
+    def _parse_names(self, names: Union[list[str], None]) -> list[str]:
         """
         Parse the given callbacks names. If they are not 'None' then the names will be returned as they are, otherwise
         all of the callbacks handled by this handler will be returned (the default behavior of when there were no names
@@ -646,7 +646,7 @@ class CallbacksHandler:
         return list(self._callbacks.keys())
 
     def _run_callbacks(
-        self, method_name: str, callbacks: List[str], *args, **kwargs
+        self, method_name: str, callbacks: list[str], *args, **kwargs
     ) -> bool:
         """
         Run the given method from the 'CallbackInterface' on all the specified callbacks with the given arguments.

mlrun/frameworks/pytorch/mlrun_interface.py

@@ -14,7 +14,7 @@
 #
 import importlib
 import sys
-from typing import Any, Dict, List, Tuple, Union
+from typing import Any, Union
 
 import torch
 import torch.multiprocessing as mp
@@ -109,13 +109,13 @@ class PyTorchMLRunInterface:
         loss_function: Module,
         optimizer: Optimizer,
         validation_set: DataLoader = None,
-        metric_functions: List[PyTorchTypes.MetricFunctionType] = None,
+        metric_functions: list[PyTorchTypes.MetricFunctionType] = None,
         scheduler=None,
         scheduler_step_frequency: Union[int, float, str] = "epoch",
         epochs: int = 1,
         training_iterations: int = None,
         validation_iterations: int = None,
-        callbacks: List[Callback] = None,
+        callbacks: list[Callback] = None,
         use_cuda: bool = True,
         use_horovod: bool = None,
     ):
@@ -221,12 +221,12 @@ class PyTorchMLRunInterface:
         self,
         dataset: DataLoader,
         loss_function: Module = None,
-        metric_functions: List[PyTorchTypes.MetricFunctionType] = None,
+        metric_functions: list[PyTorchTypes.MetricFunctionType] = None,
         iterations: int = None,
-        callbacks: List[Callback] = None,
+        callbacks: list[Callback] = None,
         use_cuda: bool = True,
         use_horovod: bool = None,
-    ) -> List[PyTorchTypes.MetricValueType]:
+    ) -> list[PyTorchTypes.MetricValueType]:
         """
         Initiate an evaluation process on this interface configuration.
 
@@ -303,9 +303,9 @@ class PyTorchMLRunInterface:
     def add_auto_logging_callbacks(
         self,
         add_mlrun_logger: bool = True,
-        mlrun_callback_kwargs: Dict[str, Any] = None,
+        mlrun_callback_kwargs: dict[str, Any] = None,
         add_tensorboard_logger: bool = True,
-        tensorboard_callback_kwargs: Dict[str, Any] = None,
+        tensorboard_callback_kwargs: dict[str, Any] = None,
     ):
         """
         Get automatic logging callbacks to both MLRun's context and Tensorboard. For further features of logging to both
@@ -347,7 +347,7 @@ class PyTorchMLRunInterface:
 
     def predict(
         self,
-        inputs: Union[Tensor, List[Tensor]],
+        inputs: Union[Tensor, list[Tensor]],
         use_cuda: bool = True,
         batch_size: int = -1,
     ) -> Tensor:
@@ -402,13 +402,13 @@ class PyTorchMLRunInterface:
         loss_function: Module = None,
         optimizer: Optimizer = None,
         validation_set: DataLoader = None,
-        metric_functions: List[PyTorchTypes.MetricFunctionType] = None,
+        metric_functions: list[PyTorchTypes.MetricFunctionType] = None,
         scheduler=None,
         scheduler_step_frequency: Union[int, float, str] = "epoch",
         epochs: int = 1,
         training_iterations: int = None,
         validation_iterations: int = None,
-        callbacks: List[Callback] = None,
+        callbacks: list[Callback] = None,
         use_cuda: bool = True,
         use_horovod: bool = None,
     ):
@@ -734,7 +734,7 @@ class PyTorchMLRunInterface:
 
     def _validate(
         self, is_evaluation: bool = False
-    ) -> Tuple[PyTorchTypes.MetricValueType, List[PyTorchTypes.MetricValueType]]:
+    ) -> tuple[PyTorchTypes.MetricValueType, list[PyTorchTypes.MetricValueType]]:
         """
         Initiate a single epoch validation.
 
@@ -817,7 +817,7 @@ class PyTorchMLRunInterface:
         )
         return loss_value, metric_values
 
-    def _print_results(self, loss_value: Tensor, metric_values: List[float]):
+    def _print_results(self, loss_value: Tensor, metric_values: list[float]):
         """
         Print the given result between each epoch.
 
@@ -832,7 +832,7 @@ class PyTorchMLRunInterface:
             + tabulate(table, headers=["Metrics", "Values"], tablefmt="pretty")
         )
 
-    def _metrics(self, y_pred: Tensor, y_true: Tensor) -> List[float]:
+    def _metrics(self, y_pred: Tensor, y_true: Tensor) -> list[float]:
         """
         Call all the metrics on the given batch's truth and prediction output.
 
@@ -860,7 +860,7 @@ class PyTorchMLRunInterface:
         average_tensor = self._hvd.allreduce(rank_value, name=name)
         return average_tensor.item()
 
-    def _get_learning_rate(self) -> Union[Tuple[str, List[Union[str, int]]], None]:
+    def _get_learning_rate(self) -> Union[tuple[str, list[Union[str, int]]], None]:
         """
         Try and get the learning rate value form the stored optimizer.
 
@@ -949,8 +949,8 @@ class PyTorchMLRunInterface:
 
     @staticmethod
     def _tensor_to_cuda(
-        tensor: Union[Tensor, Dict, List, Tuple],
-    ) -> Union[Tensor, Dict, List, Tuple]:
+        tensor: Union[Tensor, dict, list, tuple],
+    ) -> Union[Tensor, dict, list, tuple]:
         """
         Send to given tensor to cuda if it is a tensor. If the given object is a dictionary, the dictionary values will
         be sent to the function again recursively. If the given object is a list or a tuple, all the values in it will
@@ -997,7 +997,7 @@ class PyTorchMLRunInterface:
         dataset: DataLoader,
         iterations: int,
         description: str,
-        metrics: List[PyTorchTypes.MetricFunctionType],
+        metrics: list[PyTorchTypes.MetricFunctionType],
     ) -> tqdm:
         """
         Create a progress bar for training and validating / evaluating.
@@ -1028,8 +1028,8 @@ class PyTorchMLRunInterface:
     @staticmethod
     def _update_progress_bar(
         progress_bar: tqdm,
-        metrics: List[PyTorchTypes.MetricFunctionType],
-        values: List[PyTorchTypes.MetricValueType],
+        metrics: list[PyTorchTypes.MetricFunctionType],
+        values: list[PyTorchTypes.MetricValueType],
     ):
         """
         Update the progress bar metrics results.

mlrun/frameworks/pytorch/model_handler.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 #
 import os
-from typing import Dict, List, Tuple, Type, Union
+from typing import Union
 
 import numpy as np
 import torch
@@ -50,9 +50,9 @@ class PyTorchModelHandler(DLModelHandler):
         model: Module = None,
         model_path: str = None,
         model_name: str = None,
-        model_class: Union[Type[Module], str] = None,
-        modules_map: Union[Dict[str, Union[None, str, List[str]]], str] = None,
-        custom_objects_map: Union[Dict[str, Union[str, List[str]]], str] = None,
+        model_class: Union[type[Module], str] = None,
+        modules_map: Union[dict[str, Union[None, str, list[str]]], str] = None,
+        custom_objects_map: Union[dict[str, Union[str, list[str]]], str] = None,
         custom_objects_directory: str = None,
         context: mlrun.MLClientCtx = None,
         **kwargs,
@@ -136,7 +136,7 @@ class PyTorchModelHandler(DLModelHandler):
             )
 
         # Set up the base handler class:
-        super(PyTorchModelHandler, self).__init__(
+        super().__init__(
             model=model,
             model_path=model_path,
             model_name=model_name,
@@ -152,8 +152,8 @@ class PyTorchModelHandler(DLModelHandler):
 
     def set_labels(
         self,
-        to_add: Dict[str, Union[str, int, float]] = None,
-        to_remove: List[str] = None,
+        to_add: dict[str, Union[str, int, float]] = None,
+        to_remove: list[str] = None,
     ):
         """
         Update the labels dictionary of this model artifact. There are required labels that cannot be edited or removed.
@@ -162,14 +162,14 @@ class PyTorchModelHandler(DLModelHandler):
         :param to_remove: A list of labels keys to remove.
         """
         # Update the user's labels:
-        super(PyTorchModelHandler, self).set_labels(to_add=to_add, to_remove=to_remove)
+        super().set_labels(to_add=to_add, to_remove=to_remove)
 
         # Set the required labels:
         self._labels[self._LabelKeys.MODEL_CLASS_NAME] = self._model_class_name
 
     def save(
         self, output_path: str = None, **kwargs
-    ) -> Union[Dict[str, Artifact], None]:
+    ) -> Union[dict[str, Artifact], None]:
         """
         Save the handled model at the given output path.
 
@@ -182,7 +182,7 @@ class PyTorchModelHandler(DLModelHandler):
         :raise MLRunInvalidArgumentError: If an output path was not given, yet a context was not provided in
                                           initialization.
         """
-        super(PyTorchModelHandler, self).save(output_path=output_path)
+        super().save(output_path=output_path)
 
         # Set the output path:
         if output_path is None:
@@ -207,7 +207,7 @@ class PyTorchModelHandler(DLModelHandler):
 
         :raise MLRunInvalidArgumentError: If the model's class is not in the custom objects map.
         """
-        super(PyTorchModelHandler, self).load()
+        super().load()
 
         # Validate the model's class is in the custom objects map:
         if (
@@ -233,10 +233,10 @@ class PyTorchModelHandler(DLModelHandler):
     def to_onnx(
         self,
         model_name: str = None,
-        input_sample: Union[torch.Tensor, Tuple[torch.Tensor, ...]] = None,
-        input_layers_names: List[str] = None,
-        output_layers_names: List[str] = None,
-        dynamic_axes: Dict[str, Dict[int, str]] = None,
+        input_sample: Union[torch.Tensor, tuple[torch.Tensor, ...]] = None,
+        input_layers_names: list[str] = None,
+        output_layers_names: list[str] = None,
+        dynamic_axes: dict[str, dict[int, str]] = None,
         is_batched: bool = True,
         optimize: bool = True,
         output_path: str = None,
@@ -406,7 +406,7 @@ class PyTorchModelHandler(DLModelHandler):
         ]
 
         # Continue collecting from abstract class:
-        super(PyTorchModelHandler, self)._collect_files_from_store_object()
+        super()._collect_files_from_store_object()
 
     def _collect_files_from_local_path(self):
         """
@@ -443,7 +443,7 @@ class PyTorchModelHandler(DLModelHandler):
         """
         # Supported types:
         if isinstance(sample, np.ndarray):
-            return super(PyTorchModelHandler, self)._read_sample(sample=sample)
+            return super()._read_sample(sample=sample)
         elif isinstance(sample, torch.Tensor):
             return Feature(
                 value_type=PyTorchUtils.convert_torch_dtype_to_value_type(

mlrun/frameworks/pytorch/model_server.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
-from typing import Any, Dict, List, Type, Union
+from typing import Any, Union
 
 import numpy as np
 import torch
@@ -39,9 +39,9 @@ class PyTorchModelServer(V2ModelServer):
         model: Module = None,
         model_path: str = None,
         model_name: str = None,
-        model_class: Union[Type[Module], str] = None,
-        modules_map: Union[Dict[str, Union[None, str, List[str]]], str] = None,
-        custom_objects_map: Union[Dict[str, Union[str, List[str]]], str] = None,
+        model_class: Union[type[Module], str] = None,
+        modules_map: Union[dict[str, Union[None, str, list[str]]], str] = None,
+        custom_objects_map: Union[dict[str, Union[str, list[str]]], str] = None,
         custom_objects_directory: str = None,
         use_cuda: bool = True,
         to_list: bool = False,
@@ -106,7 +106,7 @@ class PyTorchModelServer(V2ModelServer):
         :param protocol:                 -
         :param class_args:               -
         """
-        super(PyTorchModelServer, self).__init__(
+        super().__init__(
             context=context,
             name=name,
             model_path=model_path,
@@ -158,7 +158,7 @@ class PyTorchModelServer(V2ModelServer):
             model=self._model_handler.model, context=self.context
         )
 
-    def predict(self, request: Dict[str, Any]) -> Union[Tensor, list]:
+    def predict(self, request: dict[str, Any]) -> Union[Tensor, list]:
         """
         Infer the inputs through the model using MLRun's PyTorch interface and return its output. The inferred data will
         be read from the "inputs" key of the request.
@@ -183,7 +183,7 @@ class PyTorchModelServer(V2ModelServer):
         # Return as list if required:
         return predictions if not self.to_list else predictions.tolist()
 
-    def explain(self, request: Dict[str, Any]) -> str:
+    def explain(self, request: dict[str, Any]) -> str:
         """
         Return a string explaining what model is being serve in this serving function and the function name.
 

mlrun/frameworks/sklearn/__init__.py

@@ -14,7 +14,7 @@
 #
 # flake8: noqa  - this is until we take care of the F401 violations with respect to __all__ & sphinx
 import warnings
-from typing import Dict, List, Union
+from typing import Union
 
 import mlrun
 from mlrun.frameworks.sklearn.metric import Metric
@@ -37,25 +37,25 @@ def apply_mlrun(
     model_name: str = "model",
     tag: str = "",
     model_path: str = None,
-    modules_map: Union[Dict[str, Union[None, str, List[str]]], str] = None,
-    custom_objects_map: Union[Dict[str, Union[str, List[str]]], str] = None,
+    modules_map: Union[dict[str, Union[None, str, list[str]]], str] = None,
+    custom_objects_map: Union[dict[str, Union[str, list[str]]], str] = None,
     custom_objects_directory: str = None,
     context: mlrun.MLClientCtx = None,
-    artifacts: Union[List[MLPlan], List[str], Dict[str, dict]] = None,
+    artifacts: Union[list[MLPlan], list[str], dict[str, dict]] = None,
     metrics: Union[
-        List[Metric],
-        List[SKLearnTypes.MetricEntryType],
-        Dict[str, SKLearnTypes.MetricEntryType],
+        list[Metric],
+        list[SKLearnTypes.MetricEntryType],
+        dict[str, SKLearnTypes.MetricEntryType],
     ] = None,
     x_test: SKLearnTypes.DatasetType = None,
     y_test: SKLearnTypes.DatasetType = None,
     sample_set: Union[SKLearnTypes.DatasetType, mlrun.DataItem, str] = None,
-    y_columns: Union[List[str], List[int]] = None,
+    y_columns: Union[list[str], list[int]] = None,
     feature_vector: str = None,
-    feature_weights: List[float] = None,
-    labels: Dict[str, Union[str, int, float]] = None,
-    parameters: Dict[str, Union[str, int, float]] = None,
-    extra_data: Dict[str, SKLearnTypes.ExtraDataType] = None,
+    feature_weights: list[float] = None,
+    labels: dict[str, Union[str, int, float]] = None,
+    parameters: dict[str, Union[str, int, float]] = None,
+    extra_data: dict[str, SKLearnTypes.ExtraDataType] = None,
     auto_log: bool = True,
     **kwargs,
 ) -> SKLearnModelHandler:
@@ -92,7 +92,7 @@ def apply_mlrun(
 
                                          {
                                              "/.../custom_model.py": "MyModel",
-                                             "/.../custom_objects.py": ["object1", "object2"]
+                                             "/.../custom_objects.py": ["object1", "object2"],
                                          }
 
                                      All the paths will be accessed from the given 'custom_objects_directory', meaning