mlrun 1.7.2rc3__py3-none-any.whl → 1.8.0__py3-none-any.whl

mlrun/frameworks/lgbm/__init__.py

@@ -11,9 +11,8 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-#
-# flake8: noqa  - this is until we take care of the F401 violations with respect to __all__ & sphinx
-from typing import Any, Union
+
+from typing import Any, Optional, Union
 
 import lightgbm as lgb
 
@@ -37,20 +36,20 @@ LGBMArtifactsLibrary = MLArtifactsLibrary
 def _apply_mlrun_on_module(
     model_name: str = "model",
     tag: str = "",
-    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,
+    modules_map: Optional[Union[dict[str, Union[None, str, list[str]]], str]] = None,
+    custom_objects_map: Optional[Union[dict[str, Union[str, list[str]]], str]] = None,
+    custom_objects_directory: Optional[str] = None,
     context: mlrun.MLClientCtx = None,
     model_format: str = LGBMModelHandler.ModelFormats.PKL,
     sample_set: Union[LGBMTypes.DatasetType, mlrun.DataItem, str] = 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, LGBMTypes.ExtraDataType] = None,
+    y_columns: Optional[Union[list[str], list[int]]] = None,
+    feature_vector: Optional[str] = None,
+    feature_weights: Optional[list[float]] = None,
+    labels: Optional[dict[str, Union[str, int, float]]] = None,
+    parameters: Optional[dict[str, Union[str, int, float]]] = None,
+    extra_data: Optional[dict[str, LGBMTypes.ExtraDataType]] = None,
     auto_log: bool = True,
-    mlrun_logging_callback_kwargs: dict[str, Any] = None,
+    mlrun_logging_callback_kwargs: Optional[dict[str, Any]] = None,
 ):
     # Apply MLRun's interface on the LightGBM module:
     LGBMMLRunInterface.add_interface(obj=lgb)
@@ -84,27 +83,29 @@ def _apply_mlrun_on_model(
     model: LGBMTypes.ModelType = None,
     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,
-    custom_objects_directory: str = None,
+    model_path: Optional[str] = None,
+    modules_map: Optional[Union[dict[str, Union[None, str, list[str]]], str]] = None,
+    custom_objects_map: Optional[Union[dict[str, Union[str, list[str]]], str]] = None,
+    custom_objects_directory: Optional[str] = None,
     context: mlrun.MLClientCtx = None,
     model_format: str = LGBMModelHandler.ModelFormats.PKL,
-    artifacts: Union[list[MLPlan], list[str], dict[str, dict]] = None,
-    metrics: Union[
-        list[Metric],
-        list[LGBMTypes.MetricEntryType],
-        dict[str, LGBMTypes.MetricEntryType],
+    artifacts: Optional[Union[list[MLPlan], list[str], dict[str, dict]]] = None,
+    metrics: Optional[
+        Union[
+            list[Metric],
+            list[LGBMTypes.MetricEntryType],
+            dict[str, LGBMTypes.MetricEntryType],
+        ]
     ] = None,
     x_test: LGBMTypes.DatasetType = None,
     y_test: LGBMTypes.DatasetType = None,
     sample_set: Union[LGBMTypes.DatasetType, mlrun.DataItem, str] = 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, LGBMTypes.ExtraDataType] = None,
+    y_columns: Optional[Union[list[str], list[int]]] = None,
+    feature_vector: Optional[str] = None,
+    feature_weights: Optional[list[float]] = None,
+    labels: Optional[dict[str, Union[str, int, float]]] = None,
+    parameters: Optional[dict[str, Union[str, int, float]]] = None,
+    extra_data: Optional[dict[str, LGBMTypes.ExtraDataType]] = None,
     auto_log: bool = True,
     **kwargs,
 ):
@@ -182,29 +183,31 @@ def apply_mlrun(
     model: LGBMTypes.ModelType = None,
     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,
-    custom_objects_directory: str = None,
+    model_path: Optional[str] = None,
+    modules_map: Optional[Union[dict[str, Union[None, str, list[str]]], str]] = None,
+    custom_objects_map: Optional[Union[dict[str, Union[str, list[str]]], str]] = None,
+    custom_objects_directory: Optional[str] = None,
     context: mlrun.MLClientCtx = None,
     model_format: str = LGBMModelHandler.ModelFormats.PKL,
-    artifacts: Union[list[MLPlan], list[str], dict[str, dict]] = None,
-    metrics: Union[
-        list[Metric],
-        list[LGBMTypes.MetricEntryType],
-        dict[str, LGBMTypes.MetricEntryType],
+    artifacts: Optional[Union[list[MLPlan], list[str], dict[str, dict]]] = None,
+    metrics: Optional[
+        Union[
+            list[Metric],
+            list[LGBMTypes.MetricEntryType],
+            dict[str, LGBMTypes.MetricEntryType],
+        ]
     ] = None,
     x_test: LGBMTypes.DatasetType = None,
     y_test: LGBMTypes.DatasetType = None,
     sample_set: Union[LGBMTypes.DatasetType, mlrun.DataItem, str] = 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, LGBMTypes.ExtraDataType] = None,
+    y_columns: Optional[Union[list[str], list[int]]] = None,
+    feature_vector: Optional[str] = None,
+    feature_weights: Optional[list[float]] = None,
+    labels: Optional[dict[str, Union[str, int, float]]] = None,
+    parameters: Optional[dict[str, Union[str, int, float]]] = None,
+    extra_data: Optional[dict[str, LGBMTypes.ExtraDataType]] = None,
     auto_log: bool = True,
-    mlrun_logging_callback_kwargs: dict[str, Any] = None,
+    mlrun_logging_callback_kwargs: Optional[dict[str, Any]] = None,
     **kwargs,
 ) -> Union[LGBMModelHandler, None]:
     """

mlrun/frameworks/lgbm/callbacks/__init__.py

@@ -11,8 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-#
-# flake8: noqa  - this is until we take care of the F401 violations with respect to __all__ & sphinx
+
 from .callback import Callback
 from .logging_callback import LoggingCallback
 from .mlrun_logging_callback import MLRunLoggingCallback

mlrun/frameworks/lgbm/callbacks/logging_callback.py

@@ -13,6 +13,8 @@
 # limitations under the License.
 #
 
+from typing import Optional
+
 from ..._ml_common.loggers import Logger
 from ..utils import LGBMTypes
 from .callback import Callback, CallbackEnv
@@ -25,8 +27,8 @@ class LoggingCallback(Callback):
 
     def __init__(
         self,
-        dynamic_hyperparameters: list[str] = None,
-        static_hyperparameters: list[str] = None,
+        dynamic_hyperparameters: Optional[list[str]] = None,
+        static_hyperparameters: Optional[list[str]] = None,
     ):
         """
         Initialize the logging callback with the given configuration. All the metrics data will be collected but the

mlrun/frameworks/lgbm/callbacks/mlrun_logging_callback.py

@@ -13,6 +13,8 @@
 # limitations under the License.
 #
 
+from typing import Optional
+
 import mlrun
 
 from ..._ml_common.loggers import MLRunLogger
@@ -33,8 +35,8 @@ class MLRunLoggingCallback(LoggingCallback):
     def __init__(
         self,
         context: mlrun.MLClientCtx,
-        dynamic_hyperparameters: list[str] = None,
-        static_hyperparameters: list[str] = None,
+        dynamic_hyperparameters: Optional[list[str]] = None,
+        static_hyperparameters: Optional[list[str]] = None,
         logging_frequency: int = 100,
     ):
         """

mlrun/frameworks/lgbm/mlrun_interfaces/__init__.py

@@ -11,8 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-#
-# flake8: noqa  - this is until we take care of the F401 violations with respect to __all__ & sphinx
+
 from .booster_mlrun_interface import LGBMBoosterMLRunInterface
 from .mlrun_interface import LGBMMLRunInterface
 from .model_mlrun_interface import LGBMModelMLRunInterface

mlrun/frameworks/lgbm/mlrun_interfaces/mlrun_interface.py

@@ -14,7 +14,7 @@
 #
 from abc import ABC
 from types import ModuleType
-from typing import Callable, Union
+from typing import Callable, Optional, Union
 
 import lightgbm as lgb
 
@@ -67,7 +67,7 @@ class LGBMMLRunInterface(MLRunInterface, ABC):
     @classmethod
     def add_interface(
         cls,
-        obj: ModuleType = None,
+        obj: Optional[ModuleType] = None,
         restoration: LGBMTypes.MLRunInterfaceRestorationType = None,
     ):
         """
@@ -167,10 +167,10 @@ class LGBMMLRunInterface(MLRunInterface, ABC):
     def configure_logging(
         context: mlrun.MLClientCtx = None,
         log_model: bool = True,
-        model_handler_kwargs: dict = None,
-        log_model_kwargs: dict = None,
+        model_handler_kwargs: Optional[dict] = None,
+        log_model_kwargs: Optional[dict] = None,
         log_training: bool = True,
-        mlrun_logging_callback_kwargs: dict = None,
+        mlrun_logging_callback_kwargs: Optional[dict] = None,
     ):
         """
         Configure the logging of the training API in LightGBM to log the training and model into MLRun. Each `train`

mlrun/frameworks/lgbm/model_handler.py

@@ -14,7 +14,7 @@
 #
 import os
 import pickle
-from typing import Union
+from typing import Optional, Union
 
 import cloudpickle
 import lightgbm as lgb
@@ -53,12 +53,16 @@ class LGBMModelHandler(MLModelHandler):
 
     def __init__(
         self,
-        model_name: str = None,
-        model_path: str = None,
+        model_name: Optional[str] = None,
+        model_path: Optional[str] = None,
         model: LGBMTypes.ModelType = 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,
+        modules_map: Optional[
+            Union[dict[str, Union[None, str, list[str]]], str]
+        ] = None,
+        custom_objects_map: Optional[
+            Union[dict[str, Union[str, list[str]]], str]
+        ] = None,
+        custom_objects_directory: Optional[str] = None,
         context: mlrun.MLClientCtx = None,
         model_format: str = ModelFormats.PKL,
         **kwargs,
@@ -152,8 +156,8 @@ class LGBMModelHandler(MLModelHandler):
 
     def set_labels(
         self,
-        to_add: dict[str, Union[str, int, float]] = None,
-        to_remove: list[str] = None,
+        to_add: Optional[dict[str, Union[str, int, float]]] = None,
+        to_remove: Optional[list[str]] = None,
     ):
         """
         Update the labels dictionary of this model artifact. There are required labels that cannot be edited or removed.
@@ -183,7 +187,7 @@ class LGBMModelHandler(MLModelHandler):
                 f"'model_path': '{self._model_path}'"
             )
 
-    def save(self, output_path: str = None, **kwargs):
+    def save(self, output_path: Optional[str] = None, **kwargs):
         """
         Save the handled model at the given output path. If a MLRun context is available, the saved model files will be
         logged and returned as artifacts.
@@ -217,10 +221,10 @@ class LGBMModelHandler(MLModelHandler):
 
     def to_onnx(
         self,
-        model_name: str = None,
+        model_name: Optional[str] = None,
         optimize: bool = True,
         input_sample: LGBMTypes.DatasetType = None,
-        log: bool = None,
+        log: Optional[bool] = None,
     ):
         """
         Convert the model in this handler to an ONNX model. The inputs names are optional, they do not change the

mlrun/frameworks/lgbm/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, Union
+from typing import Any, Optional, Union
 
 import numpy as np
 
@@ -32,16 +32,20 @@ class LGBMModelServer(V2ModelServer):
     def __init__(
         self,
         context: mlrun.MLClientCtx = None,
-        name: str = None,
+        name: Optional[str] = None,
         model: LGBMTypes.ModelType = None,
         model_path: LGBMTypes.PathType = None,
-        model_name: str = None,
+        model_name: Optional[str] = None,
         model_format: str = LGBMModelHandler.ModelFormats.PKL,
-        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,
+        modules_map: Optional[
+            Union[dict[str, Union[None, str, list[str]]], str]
+        ] = None,
+        custom_objects_map: Optional[
+            Union[dict[str, Union[str, list[str]]], str]
+        ] = None,
+        custom_objects_directory: Optional[str] = None,
         to_list: bool = True,
-        protocol: str = None,
+        protocol: Optional[str] = None,
         **class_args,
     ):
         """

mlrun/frameworks/lgbm/utils.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
-from typing import Union
+from typing import Optional, Union
 
 import lightgbm as lgb
 import numpy as np
@@ -109,7 +109,7 @@ class LGBMUtils(MLUtils):
     def get_algorithm_functionality(
         model: MLTypes.ModelType = None,
         y: MLTypes.DatasetType = None,
-        objective: str = None,
+        objective: Optional[str] = None,
     ) -> AlgorithmFunctionality:
         """
         Get the algorithm functionality of the LightGBM model. If SciKit-Learn API is used, pass the LGBBMModel and a y

mlrun/frameworks/onnx/__init__.py

@@ -11,8 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-#
-# flake8: noqa  - this is until we take care of the F401 violations with respect to __all__ & sphinx
+
 from .dataset import ONNXDataset
 from .model_handler import ONNXModelHandler
 from .model_server import ONNXModelServer

mlrun/frameworks/onnx/dataset.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 #
 import math
-from typing import Callable, Union
+from typing import Callable, Optional, Union
 
 import numpy as np
 
@@ -28,8 +28,8 @@ class ONNXDataset:
         x: Union[np.ndarray, list[np.ndarray]],
         y: Union[np.ndarray, list[np.ndarray]] = None,
         batch_size: int = 1,
-        x_transforms: list[Callable[[np.ndarray], np.ndarray]] = None,
-        y_transforms: list[Callable[[np.ndarray], np.ndarray]] = None,
+        x_transforms: Optional[list[Callable[[np.ndarray], np.ndarray]]] = None,
+        y_transforms: Optional[list[Callable[[np.ndarray], np.ndarray]]] = None,
         is_batched_transforms: bool = False,
     ):
         """

mlrun/frameworks/onnx/mlrun_interface.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
-from typing import Callable
+from typing import Callable, Optional
 
 import numpy as np
 import onnx
@@ -35,7 +35,7 @@ class ONNXMLRunInterface:
     def __init__(
         self,
         model: onnx.ModelProto,
-        execution_providers: list[str] = None,
+        execution_providers: Optional[list[str]] = None,
         context: mlrun.MLClientCtx = None,
     ):
         # Set the context:

mlrun/frameworks/onnx/model_handler.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 #
 import os
-from typing import Union
+from typing import Optional, Union
 
 import onnx
 import onnxoptimizer
@@ -35,8 +35,8 @@ class ONNXModelHandler(ModelHandler):
     def __init__(
         self,
         model: onnx.ModelProto = None,
-        model_path: str = None,
-        model_name: str = None,
+        model_path: Optional[str] = None,
+        model_name: Optional[str] = None,
         context: mlrun.MLClientCtx = None,
         **kwargs,
     ):
@@ -70,7 +70,7 @@ class ONNXModelHandler(ModelHandler):
 
     # TODO: output_path won't work well with logging artifacts. Need to look into changing the logic of 'log_artifact'.
     def save(
-        self, output_path: str = None, **kwargs
+        self, output_path: Optional[str] = None, **kwargs
     ) -> Union[dict[str, Artifact], None]:
         """
         Save the handled model at the given output path. If a MLRun context is available, the saved model files will be
@@ -106,7 +106,9 @@ class ONNXModelHandler(ModelHandler):
         # Load the ONNX model:
         self._model = onnx.load(self._model_file)
 
-    def optimize(self, optimizations: list[str] = None, fixed_point: bool = False):
+    def optimize(
+        self, optimizations: Optional[list[str]] = None, fixed_point: bool = False
+    ):
         """
         Use ONNX optimizer to optimize the ONNX model. The optimizations supported can be seen by calling
         'onnxoptimizer.get_available_passes()'

mlrun/frameworks/onnx/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, Union
+from typing import Any, Optional, Union
 
 import numpy as np
 import onnx
@@ -33,12 +33,14 @@ class ONNXModelServer(V2ModelServer):
     def __init__(
         self,
         context: mlrun.MLClientCtx = None,
-        name: str = None,
+        name: Optional[str] = None,
         model: onnx.ModelProto = None,
-        model_path: str = None,
-        model_name: str = None,
-        execution_providers: list[Union[str, tuple[str, dict[str, Any]]]] = None,
-        protocol: str = None,
+        model_path: Optional[str] = None,
+        model_name: Optional[str] = None,
+        execution_providers: Optional[
+            list[Union[str, tuple[str, dict[str, Any]]]]
+        ] = None,
+        protocol: Optional[str] = None,
         **class_args,
     ):
         """

mlrun/frameworks/parallel_coordinates.py

@@ -11,10 +11,10 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-#
-import datetime
+
 import os
-from typing import Union
+from datetime import datetime
+from typing import Optional, Union
 
 import numpy as np
 import pandas as pd
@@ -48,7 +48,7 @@ def _gen_dropdown_buttons(output_cols) -> list:
 
 
 def _gen_dimensions(
-    df: pd.DataFrame, col: str, prefix: str = None, is_index=False
+    df: pd.DataFrame, col: str, prefix: Optional[str] = None, is_index=False
 ) -> dict:
     """
     Computes the plotting dimensions of each parameter/output col according to its type.
@@ -107,8 +107,8 @@ def gen_pcp_plot(
     source_df: pd.DataFrame,
     index_col: str,
     hide_identical: bool = True,
-    exclude: list = None,
-    colorscale: str = None,
+    exclude: Optional[list] = None,
+    colorscale: Optional[str] = None,
 ):
     """
     Creates a list composed of the data to be plotted as a Parallel Coordinate, this includes
@@ -240,11 +240,11 @@ def _runs_list_to_df(runs_list, extend_iterations=False):
 def compare_run_objects(
     runs_list: Union[mlrun.model.RunObject, list[mlrun.model.RunObject]],
     hide_identical: bool = True,
-    exclude: list = None,
-    show: bool = None,
+    exclude: Optional[list] = None,
+    show: Optional[bool] = None,
     extend_iterations=True,
     filename=None,
-    colorscale: str = None,
+    colorscale: Optional[str] = None,
 ):
     """return/show parallel coordinates plot + table to compare between a list of runs or run iterations
 
@@ -292,9 +292,9 @@ def compare_db_runs(
     run_name=None,
     labels=None,
     iter=False,
-    start_time_from: datetime = None,
+    start_time_from: Optional[datetime] = None,
     hide_identical: bool = True,
-    exclude: list = None,
+    exclude: Optional[list] = None,
     show=None,
     colorscale: str = "Blues",
     filename=None,

mlrun/frameworks/pytorch/__init__.py

@@ -11,9 +11,8 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-#
-# flake8: noqa  - this is until we take care of the F401 violations with respect to __all__ & sphinx
-from typing import Any, Union
+
+from typing import Any, Optional, Union
 
 from torch.nn import Module
 from torch.optim import Optimizer
@@ -35,23 +34,23 @@ def train(
     loss_function: Module,
     optimizer: Optimizer,
     validation_set: DataLoader = None,
-    metric_functions: list[PyTorchTypes.MetricFunctionType] = None,
+    metric_functions: Optional[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: list[Callback] = None,
+    training_iterations: Optional[int] = None,
+    validation_iterations: Optional[int] = None,
+    callbacks_list: Optional[list[Callback]] = None,
     use_cuda: bool = True,
-    use_horovod: bool = None,
+    use_horovod: Optional[bool] = None,
     auto_log: bool = True,
-    model_name: 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,
-    tensorboard_directory: str = None,
-    mlrun_callback_kwargs: dict[str, Any] = None,
-    tensorboard_callback_kwargs: dict[str, Any] = None,
+    model_name: Optional[str] = None,
+    modules_map: Optional[Union[dict[str, Union[None, str, list[str]]], str]] = None,
+    custom_objects_map: Optional[Union[dict[str, Union[str, list[str]]], str]] = None,
+    custom_objects_directory: Optional[str] = None,
+    tensorboard_directory: Optional[str] = None,
+    mlrun_callback_kwargs: Optional[dict[str, Any]] = None,
+    tensorboard_callback_kwargs: Optional[dict[str, Any]] = None,
     context: mlrun.MLClientCtx = None,
 ) -> PyTorchModelHandler:
     """
@@ -205,17 +204,17 @@ def evaluate(
     dataset: DataLoader,
     model: Module = None,
     loss_function: Module = None,
-    metric_functions: list[PyTorchTypes.MetricFunctionType] = None,
-    iterations: int = None,
-    callbacks_list: list[Callback] = None,
+    metric_functions: Optional[list[PyTorchTypes.MetricFunctionType]] = None,
+    iterations: Optional[int] = None,
+    callbacks_list: Optional[list[Callback]] = None,
     use_cuda: bool = True,
     use_horovod: bool = False,
     auto_log: bool = True,
-    model_name: 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,
-    mlrun_callback_kwargs: dict[str, Any] = None,
+    model_name: Optional[str] = None,
+    modules_map: Optional[Union[dict[str, Union[None, str, list[str]]], str]] = None,
+    custom_objects_map: Optional[Union[dict[str, Union[str, list[str]]], str]] = None,
+    custom_objects_directory: Optional[str] = None,
+    mlrun_callback_kwargs: Optional[dict[str, Any]] = None,
     context: mlrun.MLClientCtx = None,
 ) -> tuple[PyTorchModelHandler, list[PyTorchTypes.MetricValueType]]:
     """

mlrun/frameworks/pytorch/callbacks/__init__.py

@@ -11,8 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-#
-# flake8: noqa  - this is until we take care of the F401 violations with respect to __all__ & sphinx
+
 from .callback import Callback
 from .logging_callback import HyperparametersKeys, LoggingCallback
 from .mlrun_logging_callback import MLRunLoggingCallback

mlrun/frameworks/pytorch/callbacks/callback.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 #
 from abc import ABC, abstractmethod
+from typing import Optional
 
 from torch import Tensor
 from torch.nn import Module
@@ -67,7 +68,7 @@ class Callback(ABC):
         validation_set: DataLoader = None,
         loss_function: Module = None,
         optimizer: Optimizer = None,
-        metric_functions: list[PyTorchTypes.MetricFunctionType] = None,
+        metric_functions: Optional[list[PyTorchTypes.MetricFunctionType]] = None,
         scheduler=None,
     ):
         """