docstrings and typehints

Author: eaidova

optimum/exporters/openvino/convert.py

@@ -25,8 +25,9 @@
 from openvino.runtime.utils.types import get_element_type
 from openvino.tools.ovc import convert_model
 from optimum.exporters.onnx.base import OnnxConfig
-from optimum.exporters.onnx.convert import check_dummy_inputs_are_allowed, export_tensorflow as export_tensorflow_onnx
+from optimum.exporters.onnx.convert import check_dummy_inputs_are_allowed
 from optimum.exporters.onnx.convert import export_pytorch as export_pytorch_to_onnx
+from optimum.exporters.onnx.convert import export_tensorflow as export_tensorflow_onnx
 from optimum.utils import is_diffusers_available
 
 from ...intel.openvino.utils import OV_XML_FILE_NAME
@@ -119,6 +120,20 @@ def export(
 
 
 def export_tensorflow(model: Union["PreTrainedModel", "ModelMixin"], config: OnnxConfig, opset: int, output: Path):
+    """
+    Export the TensorFlow model to OpenVINO format.
+
+    Args:
+        model (Union[): The model to export.
+        config (OnnxConfig): The configuration of the model.
+        opset (int): The ONNX opset version to use.
+        output (Path): The path to save the model.
+
+    Returns:
+        input_names: list of input names from ONNX configuration
+        output_names: list of output names from ONNX configuration
+        bool:  True if the model was exported successfully.
+    """
     onnx_path = Path(output).with_suffix(".onnx")
     input_names, output_names = export_tensorflow_onnx(model, config, opset, onnx_path)
     ov_model = convert_model(str(onnx_path))
@@ -139,6 +154,30 @@ def export_pytorch_via_onnx(
     input_shapes: Optional[Dict] = None,
     model_kwargs: Optional[Dict[str, Any]] = None,
 ):
+    """
+    Exports a PyTorch model to an OpenVINO Intermediate Representation via ONNX export.
+
+    Args:
+        model ([`PreTrainedModel`]):
+            The model to export.
+        config ([`~exporters.onnx.config.OnnxConfig`]):
+            The configuration associated with the exported model.
+        opset (`int`):
+            The version of the ONNX operator set to use.
+        output (`Path`):
+            Directory to store the exported model.
+        device (`str`, defaults to `"cpu"`):
+            The device on which the model will be exported. Either `cpu` or `cuda`. Only PyTorch is supported for
+            export on CUDA devices.
+        input_shapes (`optional[Dict]`, defaults to `None`):
+            If specified, allows to use specific shapes for the example input provided to the exporter.
+        model_kwargs (optional[Dict[str, Any]], defaults to `None`):
+            Additional kwargs for model export
+
+    Returns:
+        `Tuple[List[str], List[str], bool]`: A tuple with an ordered list of the model's inputs, and the named inputs from
+        the ONNX configuration and boolean flag - was legacy ONNX path were applied to model or not.
+    """
     import torch
 
     output = Path(output)
@@ -186,10 +225,12 @@ def export_pytorch(
             export on CUDA devices.
         input_shapes (`optional[Dict]`, defaults to `None`):
             If specified, allows to use specific shapes for the example input provided to the exporter.
+        model_kwargs (optional[Dict[str, Any]], defaults to `None`):
+            Additional kwargs for model export
 
     Returns:
-        `Tuple[List[str], List[str]]`: A tuple with an ordered list of the model's inputs, and the named inputs from
-        the ONNX configuration.
+        `Tuple[List[str], List[str], bool]`: A tuple with an ordered list of the model's inputs, and the named inputs from
+        the ONNX configuration and boolean flag - was legacy ONNX path were applied to model or not.
     """
     import torch
     from torch.utils._pytree import tree_map
@@ -299,6 +340,28 @@ def export_models(
     input_shapes: Optional[Dict] = None,
     model_kwargs: Optional[Dict[str, Any]] = None,
 ) -> Tuple[List[List[str]], List[List[str]]]:
+    """
+    Export the models to OpenVINO IR format
+
+    Args:
+        models_and_onnx_configs (Dict[ str, Tuple[Union["PreTrainedModel", "TFPreTrainedModel", "ModelMixin"], "OnnxConfig"]):
+        output_dir (Path): output directory for saving models
+        opset (Optional[int], optional, Default to None): ONNX export opset
+        output_names (Optional[List[str]], optional, Defaults to None): model output names
+        device (str, optional, Defaults to "cpu"):
+            The device on which the model will be exported. Either `cpu` or `cuda`. Only PyTorch is supported for
+            export on CUDA devices.
+        input_shapes (Optional[Dict], optional, Defaults to None):
+            If specified, allows to use specific shapes for the example input provided to the exporter.
+        model_kwargs (Optional[Dict[str, Any]], optional):
+            Additional kwargs for model export
+
+    Raises:
+        ValueError: if custom names set not equal of number of models
+
+    Returns:
+        list of input_names and output_names from ONNX configuration
+    """
     outputs = []
 
     if output_names is not None and len(output_names) != len(models_and_onnx_configs):

optimum/exporters/openvino/utils.py

@@ -12,25 +12,48 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 
+from typing import Any, Dict, List, Tuple, Union
+
 from transformers.utils import is_torch_available
 
 from openvino.runtime import PartialShape
-
-from ...intel.utils.import_utils import is_nncf_available
+from optimum.utils import is_diffusers_available
 
 
 if is_torch_available():
     import torch
     import torch.nn as nn
+    from transformers.modeling_utils import PreTrainedModel
+
+if is_diffusers_available():
+    from diffusers import ModelMixin
+
 
+def is_torch_model(model: Union["PreTrainedModel", "ModelMixin"]):
+    """
+    Checks whether the model is a torch model.
 
-def is_torch_model(model):
+    Args:
+        model (Union[PretrainedModel, ModelMixin]): The model to check.
+
+    Returns:
+        bool: True if the model is a torch model.
+    """
     if not is_torch_available():
         return False
     return isinstance(model, nn.Module)
 
 
-def flattenize_inputs(inputs):
+def flattenize_inputs(inputs: List[Any]):
+    """
+    Flatten the inputs into a list.
+
+    Args:
+        inputs (List[Any]): The inputs to flatten.
+
+    Returns:
+        List[Any]:  The flattened inputs.
+    """
     flatten_inputs = []
     for input_data in inputs:
         if input_data is None:
@@ -42,8 +65,27 @@ def flattenize_inputs(inputs):
     return flatten_inputs
 
 
-def remove_none_from_dummy_inputs(dummy_inputs):
-    def remove_none_from_list_tuple(item):
+def remove_none_from_dummy_inputs(dummy_inputs: Dict[str, Any]):
+    """
+    Removes None values from the dictionary.
+
+    Args:
+        dummy_inputs (Dict[str, Any]): Dictionary with None values.
+    Returns:
+        upd_dummy (Dict[str, Any]): updated dictionary with removed None values
+        dict_dummy (List[Tuple[str, List[str]]]): list of inputs represented as dictionary provided as pair name and list of nested keys
+    """
+
+    def remove_none_from_list_tuple(item: Union[List[Any], Tuple[Any]]):
+        """
+        Removes None values from a list or tuple.
+
+        Args:
+            item (list or tuple): The list or tuple to remove None values from.
+
+        Returns:
+            list or tuple: The list or tuple with None values removed.
+        """
         new_item = [i for i in item if i is not None]
         return type(item)(new_item)
 
@@ -63,7 +105,18 @@ def remove_none_from_list_tuple(item):
     return upd_dummy, dict_dummy
 
 
-def get_input_shapes(dummy_inputs, inputs):
+def get_input_shapes(dummy_inputs: Dict[str, Any], inputs: Dict[str, Any]):
+    """
+    Resolves input shapes based on dynamic axes from input config and dummy input shapes
+
+    Args:
+        dummy_inputs (Dict[str, Any]): A dictionary of dummy inputs.
+        inputs (Dict[str, Any]): A dictionary of input tensors.
+
+    Returns:
+       input_info: List of input info for conversion
+
+    """
     input_info = []
     for input_name, data in dummy_inputs.items():
         if isinstance(data, (tuple, list, dict)):
@@ -78,6 +131,9 @@ def get_input_shapes(dummy_inputs, inputs):
 
 
 def clear_class_registry():
+    """
+    Removes Torchscript cached modules
+    """
     torch._C._jit_clear_class_registry()
     torch.jit._recursive.concrete_type_store = torch.jit._recursive.ConcreteTypeStore()
     torch.jit._state._clear_class_state()

optimum/intel/openvino/modeling_base.py

@@ -25,8 +25,8 @@
 from transformers import PretrainedConfig
 from transformers.file_utils import add_start_docstrings
 
-from optimum.exporters.tasks import TasksManager
 from optimum.exporters.onnx.base import OnnxConfig
+from optimum.exporters.tasks import TasksManager
 from optimum.modeling_base import OptimizedModel
 
 from ...exporters.openvino import export

optimum/intel/utils/modeling_utils.py

@@ -15,6 +15,7 @@
 from typing import Tuple
 
 import torch
+from transformers.modeling_utils import PreTrainedModel
 
 
 # Modified from transformers.models.bloom.modeling_bloom._make_causal_mask
@@ -91,7 +92,16 @@ def _prepare_decoder_attention_mask(attention_mask, input_shape, inputs_embeds,
     return combined_attention_mask
 
 
-def patch_decoder_attention_mask(model):
+def patch_decoder_attention_mask(model: "PreTrainedModel"):
+    """
+    Apply patch on decoder with past model forward to resolve first inference based on model architecture
+
+    Args:
+        model (PretrainedModel): The model to patch.
+
+    Returns:
+        model with applied patch
+    """
     if model.config.model_type == "bloom":
         model.transformer._prepare_attn_mask = _prepare_attn_mask
     elif model.config.model_type == "llama":

setup.py

@@ -42,7 +42,7 @@
         "onnxruntime<1.15.0",
     ],
     "openvino": ["openvino==2023.1.0.dev20230811", "onnx", "onnxruntime"],
-    "nncf": ["nncf @ git+https://github.com/openvinotoolkit/nncf.git"],
+    "nncf": ["nncf @ git+https://github.com/openvinotoolkit/nncf.git", "transformers<4.32.0"],
     "ipex": ["transformers<4.32.0", "intel-extension-for-pytorch", "onnx"],
     "diffusers": ["diffusers", "invisible-watermark>=0.2.0"],
     "quality": QUALITY_REQUIRE,