leap-model-parser 0.1.225__tar.gz → 0.1.225.dev2__tar.gz

{leap_model_parser-0.1.225 → leap_model_parser-0.1.225.dev2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: leap-model-parser
-Version: 0.1.225
+Version: 0.1.225.dev2
 Summary: 
 Home-page: https://github.com/tensorleap/leap-model-parser
 License: MIT

{leap_model_parser-0.1.225 → leap_model_parser-0.1.225.dev2}/leap_model_parser/contract/graph.py

@@ -47,7 +47,7 @@ class Node:
     outputs: Dict[str, ConnectionOutput] = field(default_factory=dict)
     pruning_plan_id: Optional[str] = None
     wrapper: Optional[WrapperData] = None
-    shape: Optional[List[str]] = None
+    shape: Optional[List[Optional[str]]] = None
 
     def __key(self):
         return (self.id, self.name)

{leap_model_parser-0.1.225 → leap_model_parser-0.1.225.dev2}/leap_model_parser/keras_json_model_import.py

@@ -205,7 +205,7 @@ class KerasJsonModelImport:
                            'serialized_call_kwargs': serialized_call_kwargs})
 
     def generate_regular_node(self, layer: Dict[str, Any], layer_metadata: Dict[str, Any], node_key: str,
-                              layer_name_to_inbound_nodes: Dict[str, List[keras_node]]):
+                              layer_name_to_inbound_nodes: Dict[str, List[Node]]):
         data = layer['config']
         if layer['class_name'] in ('TFOpLambda', 'SlicingOpLambda') or layer['class_name'] in self.custom_layers:
             call_args = layer_name_to_inbound_nodes[layer['config']

{leap_model_parser-0.1.225 → leap_model_parser-0.1.225.dev2}/leap_model_parser/model_parser.py

@@ -7,6 +7,7 @@ import tempfile
 from importlib.util import find_spec
 from pathlib import Path
 
+
 import tensorflow as tf  # type: ignore
 from code_loader.contract.mapping import NodeConnection, NodeMapping  # type: ignore
 from keras import Model  # type: ignore
@@ -33,7 +34,6 @@ if spec is not None:
 
     onnx_imported = True
 
-
 class ModelParser:
     def __init__(self, should_transform_inputs_and_outputs=False,
                  custom_layers=None,
@@ -54,14 +54,12 @@ class ModelParser:
             ImportModelTypeEnum.PB_TF2.value: self.convert_pb_model,
         }
 
-    @staticmethod
-    def _add_output_node_shape_to_model_schema(model_schema: Dict, keras_model: Model):
+    def _add_output_node_shape_to_model_schema(self, model_schema: Dict, keras_model: Model):
         for i, layer in enumerate(keras_model.layers):
             model_schema['config']['layers'][i]['output_shape'] = list(layer.output_shape)
 
     def get_keras_model_and_model_graph(
-            self, model_path: Path, model_type: ImportModelTypeEnum) -> Tuple[
-        Dict[str, Node], List[InputInfo], Optional[Model], Optional[str]]:
+            self, model_path: Path, model_type: ImportModelTypeEnum) -> Tuple[Dict[str, Node], List[InputInfo], Optional[Model], Optional[str]]:
         model_to_keras_converter: Optional[Callable[[str], Tuple[Dict[str, Node], Model, Optional[str]]]] = \
             self._model_types_converter.get(model_type.value)
         if model_to_keras_converter is None:
@@ -92,6 +90,7 @@ class ModelParser:
 
         return graph, connected_inputs, keras_model_with_weights, error_info
 
+
     def _get_k_model_from_pb_path(self, file_path: str):
         tar_file = tarfile.open(file_path)
         with tempfile.TemporaryDirectory() as temp_dir:
@@ -106,8 +105,7 @@ class ModelParser:
             k_model = self._load_keras_model_with_custom_layers(pb_folder_path)
         return k_model
 
-    def generate_model_graph(self, model_path: Path, model_type: ImportModelTypeEnum) -> Tuple[
-        Dict[str, Node], List[InputInfo]]:
+    def generate_model_graph(self, model_path: Path, model_type: ImportModelTypeEnum) -> Tuple[Dict[str, Node], List[InputInfo]]:
         model_graph, connected_inputs, _, error_info = self.get_keras_model_and_model_graph(
             model_path, model_type)
         return model_graph, connected_inputs
@@ -132,30 +130,20 @@ class ModelParser:
             node.name for node in onnx_model.graph.initializer]
         input_names = list(set(input_all) - set(input_initializer))
         converted_response: ConvertedResponse = onnx_to_keras(onnx_model, input_names=input_names,
-                                                              name_policy='attach_weights_name',
-                                                              allow_partial_compilation=False)
+                                name_policy='attach_weights_name', allow_partial_compilation=False)
         return self.convert_to_keras_model(converted_response.converted_model, converted_response.error_info)
 
     def _load_keras_model_with_custom_layers(self, file_path: str):
         custom_objects = {}
         if self.custom_layers is not None:
             custom_objects = self.custom_layers
-
-        try:
-            return load_model(file_path, custom_objects=custom_objects, compile=False)
-        except OSError as e:
-            if 'signature' in str(e):
-                raise Exception('Unable to open model file. The model might be corrupted or not a valid.')
-            else:
-                raise e
-
+        return load_model(file_path, custom_objects=custom_objects, compile=False)
 
     def convert_h5_model(self, file_path: str) -> Tuple[Dict[str, Node], Model, Optional[str]]:
         imported_model = self._load_keras_model_with_custom_layers(file_path)
         return self.convert_to_keras_model(imported_model)
 
-    def convert_to_keras_model(self, k_model, error_info: Optional[str] = None) -> Tuple[
-        Dict[str, Node], Model, Optional[str]]:
+    def convert_to_keras_model(self, k_model, error_info: Optional[str] = None) -> Tuple[Dict[str, Node], Model, Optional[str]]:
         converted_k_model = convert_channels_first_to_last(
             k_model, self._should_transform_inputs_and_outputs, self.custom_layers)
 

{leap_model_parser-0.1.225 → leap_model_parser-0.1.225.dev2}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "leap-model-parser"
-version = "0.1.225"
+version = "0.1.225.dev2"
 description = ""
 authors = ["idan <idan.yogev@tensorleap.ai>"]
 license = "MIT"