mlrun 1.7.1rc4__py3-none-any.whl → 1.8.0rc8__py3-none-any.whl

mlrun/frameworks/_dl_common/loggers/mlrun_logger.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 #
 import itertools
-from typing import Union
+from typing import Optional, Union
 
 import numpy as np
 import plotly.graph_objects as go
@@ -144,9 +144,9 @@ class MLRunLogger(Logger):
         self,
         model_handler: DLModelHandler,
         tag: str = "",
-        labels: dict[str, DLTypes.TrackableType] = None,
-        parameters: dict[str, DLTypes.TrackableType] = None,
-        extra_data: dict[str, Union[DLTypes.TrackableType, Artifact]] = None,
+        labels: Optional[dict[str, DLTypes.TrackableType]] = None,
+        parameters: Optional[dict[str, DLTypes.TrackableType]] = None,
+        extra_data: Optional[dict[str, Union[DLTypes.TrackableType, Artifact]]] = None,
     ):
         """
         Log the run, summarizing the validation metrics and dynamic hyperparameters across all epochs. If 'update' is

mlrun/frameworks/_dl_common/loggers/tensorboard_logger.py

@@ -15,7 +15,7 @@
 import os
 from abc import abstractmethod
 from datetime import datetime
-from typing import Any, Callable, Generic, Union
+from typing import Any, Callable, Generic, Optional, Union
 
 import yaml
 
@@ -64,8 +64,8 @@ class TensorboardLogger(Logger, Generic[DLTypes.WeightType]):
             Callable[[DLTypes.WeightType], Union[float, DLTypes.WeightType]]
         ],
         context: mlrun.MLClientCtx = None,
-        tensorboard_directory: str = None,
-        run_name: str = None,
+        tensorboard_directory: Optional[str] = None,
+        run_name: Optional[str] = None,
         update_frequency: Union[int, str] = "epoch",
     ):
         """

mlrun/frameworks/_ml_common/__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 .artifacts_library import MLArtifactsLibrary
 from .model_handler import MLModelHandler
 from .pkl_model_server import PickleModelServer

mlrun/frameworks/_ml_common/loggers/__init__.py

@@ -11,7 +11,6 @@
 # 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 .logger import Logger
 from .mlrun_logger import MLRunLogger

mlrun/frameworks/_ml_common/model_handler.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 #
 from abc import ABC
-from typing import Union
+from typing import Optional, Union
 
 import mlrun
 from mlrun.artifacts import Artifact
@@ -33,7 +33,7 @@ class MLModelHandler(ModelHandler, ABC):
         self,
         model: MLTypes.ModelType = None,
         model_path: MLTypes.PathType = None,
-        model_name: str = None,
+        model_name: Optional[str] = None,
         modules_map: Union[
             dict[str, Union[None, str, list[str]]], MLTypes.PathType
         ] = None,
@@ -224,18 +224,18 @@ class MLModelHandler(ModelHandler, ABC):
     def log(
         self,
         tag: str = "",
-        labels: dict[str, Union[str, int, float]] = None,
-        parameters: dict[str, Union[str, int, float]] = None,
-        inputs: list[Feature] = None,
-        outputs: list[Feature] = None,
-        metrics: dict[str, Union[int, float]] = None,
-        artifacts: dict[str, Artifact] = None,
-        extra_data: dict[str, MLTypes.ExtraDataType] = None,
-        algorithm: str = None,
+        labels: Optional[dict[str, Union[str, int, float]]] = None,
+        parameters: Optional[dict[str, Union[str, int, float]]] = None,
+        inputs: Optional[list[Feature]] = None,
+        outputs: Optional[list[Feature]] = None,
+        metrics: Optional[dict[str, Union[int, float]]] = None,
+        artifacts: Optional[dict[str, Artifact]] = None,
+        extra_data: Optional[dict[str, MLTypes.ExtraDataType]] = None,
+        algorithm: Optional[str] = None,
         sample_set: MLTypes.DatasetType = None,
         target_columns: MLTypes.TargetColumnsNamesType = None,
-        feature_vector: str = None,
-        feature_weights: list[float] = None,
+        feature_vector: Optional[str] = None,
+        feature_weights: Optional[list[float]] = None,
     ):
         """
         Log the model held by this handler into the MLRun context provided.
@@ -299,15 +299,15 @@ class MLModelHandler(ModelHandler, ABC):
 
     def update(
         self,
-        labels: dict[str, Union[str, int, float]] = None,
-        parameters: dict[str, Union[str, int, float]] = None,
-        inputs: list[Feature] = None,
-        outputs: list[Feature] = None,
-        metrics: dict[str, Union[int, float]] = None,
-        artifacts: dict[str, Artifact] = None,
-        extra_data: dict[str, MLTypes.ExtraDataType] = None,
-        feature_vector: str = None,
-        feature_weights: list[float] = None,
+        labels: Optional[dict[str, Union[str, int, float]]] = None,
+        parameters: Optional[dict[str, Union[str, int, float]]] = None,
+        inputs: Optional[list[Feature]] = None,
+        outputs: Optional[list[Feature]] = None,
+        metrics: Optional[dict[str, Union[int, float]]] = None,
+        artifacts: Optional[dict[str, Artifact]] = None,
+        extra_data: Optional[dict[str, MLTypes.ExtraDataType]] = None,
+        feature_vector: Optional[str] = None,
+        feature_weights: Optional[list[float]] = None,
     ):
         """
         Update the model held by this handler into the MLRun context provided, updating the model's artifact properties

mlrun/frameworks/_ml_common/plans/__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 .calibration_curve_plan import CalibrationCurvePlan
 from .confusion_matrix_plan import ConfusionMatrixPlan
 from .dataset_plan import DatasetPlan

mlrun/frameworks/_ml_common/plans/confusion_matrix_plan.py

@@ -13,6 +13,8 @@
 # limitations under the License.
 #
 
+from typing import Optional
+
 import numpy as np
 import pandas as pd
 from plotly.figure_factory import create_annotated_heatmap
@@ -35,7 +37,7 @@ class ConfusionMatrixPlan(MLPlotPlan):
         self,
         labels: np.ndarray = None,
         sample_weight: np.ndarray = None,
-        normalize: str = None,
+        normalize: Optional[str] = None,
     ):
         """
         Initialize a confusion matrix plan with the given configuration.

mlrun/frameworks/_ml_common/plans/dataset_plan.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 #
 from enum import Enum
-from typing import Union
+from typing import Optional, Union
 
 import mlrun.errors
 from mlrun.artifacts import Artifact, DatasetArtifact
@@ -49,8 +49,8 @@ class DatasetPlan(MLPlan):
     def __init__(
         self,
         purpose: Union[Purposes, str] = Purposes.OTHER,
-        name: str = None,
-        preview: int = None,
+        name: Optional[str] = None,
+        preview: Optional[int] = None,
         stats: bool = False,
         fmt: str = "parquet",
     ):

mlrun/frameworks/_ml_common/plans/roc_curve_plan.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 numpy as np
 import pandas as pd
@@ -35,13 +35,13 @@ class ROCCurvePlan(MLPlotPlan):
 
     def __init__(
         self,
-        pos_label: Union[str, int] = None,
+        pos_label: Optional[Union[str, int]] = None,
         sample_weight: np.ndarray = None,
         drop_intermediate: bool = True,
         average: str = "macro",
-        max_fpr: float = None,
+        max_fpr: Optional[float] = None,
         multi_class: str = "raise",
-        labels: list[str] = None,
+        labels: Optional[list[str]] = None,
     ):
         """
         Initialize a receiver operating characteristic plan with the given configuration.

mlrun/frameworks/auto_mlrun/__init__.py

@@ -11,6 +11,5 @@
 # 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 .auto_mlrun import AutoMLRun

mlrun/frameworks/auto_mlrun/auto_mlrun.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 Callable, Union
+
+from typing import Callable, Optional, Union
 
 import mlrun
 from mlrun.artifacts import get_model
@@ -261,7 +260,7 @@ class AutoMLRun:
 
     @staticmethod
     def _get_framework(
-        model: CommonTypes.ModelType = None, model_path: str = None
+        model: CommonTypes.ModelType = None, model_path: Optional[str] = None
     ) -> Union[tuple[str, dict]]:
         """
         Try to get the framework from the model or model path provided. The framework can be read from the model path
@@ -320,12 +319,16 @@ class AutoMLRun:
     @staticmethod
     def load_model(
         model_path: str,
-        model_name: str = None,
+        model_name: Optional[str] = None,
         context: mlrun.MLClientCtx = 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,
-        framework: 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,
+        framework: Optional[str] = None,
         **kwargs,
     ) -> ModelHandler:
         """
@@ -417,14 +420,18 @@ class AutoMLRun:
     @staticmethod
     def apply_mlrun(
         model: CommonTypes.ModelType = None,
-        model_name: str = None,
+        model_name: Optional[str] = None,
         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,
-        framework: str = None,
+        framework: Optional[str] = None,
         auto_log: bool = True,
         **kwargs,
     ) -> ModelHandler:

mlrun/frameworks/huggingface/__init__.py

@@ -11,6 +11,5 @@
 # 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 .model_server import HuggingFaceModelServer

mlrun/frameworks/huggingface/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
+from typing import Any, Optional
 
 import numpy as np
 import transformers
@@ -32,14 +32,14 @@ class HuggingFaceModelServer(V2ModelServer):
     def __init__(
         self,
         context: mlrun.MLClientCtx = None,
-        name: str = None,
-        task: str = None,
-        model_path: str = None,
-        model_name: str = None,
-        model_class: str = None,
-        tokenizer_name: str = None,
-        tokenizer_class: str = None,
-        framework: str = None,
+        name: Optional[str] = None,
+        task: Optional[str] = None,
+        model_path: Optional[str] = None,
+        model_name: Optional[str] = None,
+        model_class: Optional[str] = None,
+        tokenizer_name: Optional[str] = None,
+        tokenizer_class: Optional[str] = None,
+        framework: Optional[str] = None,
         **class_args,
     ):
         """

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