leap-model-parser 0.1.185.dev12__tar.gz → 0.1.187__tar.gz

{leap_model_parser-0.1.185.dev12 → leap_model_parser-0.1.187}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: leap-model-parser
-Version: 0.1.185.dev12
+Version: 0.1.187
 Summary: 
 Home-page: https://github.com/tensorleap/leap-model-parser
 License: MIT
@@ -12,7 +12,7 @@ Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.8
 Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
-Requires-Dist: code-loader (==1.0.87.dev8)
+Requires-Dist: code-loader (==1.0.89)
 Requires-Dist: keras-data-format-converter (==0.1.22)
 Requires-Dist: leap-model-rebuilder (==0.1.7)
 Requires-Dist: numpy (>=1.22.3,<2.0.0)

{leap_model_parser-0.1.185.dev12 → leap_model_parser-0.1.187}/leap_model_parser/leap_graph_editor.py

@@ -1,8 +1,8 @@
 from enum import Enum
 from typing import Optional, Dict, Any, List
 
-from code_loader.contract.mapping import NodeConnection, NodeMappingType, NodeMapping
-from keras import Model
+from code_loader.contract.mapping import NodeConnection, NodeMappingType, NodeMapping  # type: ignore
+from keras import Model  # type: ignore
 
 from leap_model_parser.contract.graph import Node as Node, OutputData, ConnectionOutput, ConnectionInput, InputData
 
@@ -26,7 +26,10 @@ class LeapGraphEditor:
         if node_connection.node.type == NodeMappingType.PredictionLabels:
             prediction_labels_name = node_connection.node.name
             prediction_node = list(node_connection.node_inputs.values())[0]
-            prediction_node_id = self._find_node_by_origin_name(prediction_node.name).id
+            prediction_mapping_node = self._find_node_by_origin_name(prediction_node.name)
+            assert prediction_mapping_node is not None, \
+                f"Prediction node with name {prediction_node.name} not found in model graph"
+            prediction_node_id = prediction_mapping_node.id
             self.model_graph[prediction_node_id].data['prediction_type'] = prediction_labels_name
         if node_connection.node.type == NodeMappingType.Visualizer:
             new_visualizer_node_id = self._add_visualizer_node(
@@ -198,12 +201,15 @@ class LeapGraphEditor:
         input_name = f'{node_id}-{input_name}'
         self.model_graph[node_id].inputs[input_name] = ConnectionInput([InputData(input_node_id, output_name)])
 
+        output_connection = OutputData(node_id, input_name)
         self.model_graph[input_node_id].outputs[output_name].connections.append(output_connection)
 
-    def _handle_input_node_with_index(self, input_node: NodeMapping):
+    def _handle_input_node_with_index(self, input_node: NodeMapping) -> str:
         input_index = int(input_node.type.value.replace('Input', ''))
         origin_name = self.keras_model.inputs[input_index].node.layer.name
-        input_node_id = self._find_input_node_by_origin_name(origin_name).id
+        input_node = self._find_input_node_by_origin_name(origin_name)
+        assert input_node is not None, f"Input node with origin name {origin_name} not found in model graph"
+        input_node_id = input_node.id
         self.model_graph[input_node_id].data['output_name'] = input_node.name
         output_keys = list(self.model_graph[input_node_id].outputs.keys())
         for output_key in output_keys:
@@ -218,9 +224,6 @@ class LeapGraphEditor:
                     if connection_input.output == output_key:
                         connection_input.output = new_output_key
 
-        if input_node_id is None:
-            raise Exception(f"Couldn't find input node by origin name {origin_name}")
-
         return input_node_id
 
     def _find_or_add_input_node(self, input_node: NodeMapping) -> str:
@@ -234,7 +237,10 @@ class LeapGraphEditor:
         elif 'Input' in input_node.type.value:
             input_node_id = self._handle_input_node_with_index(input_node)
         elif input_node.type.value.startswith('Prediction'):
-            input_node_id = self._find_node_by_origin_name(input_node.name).id
+            node_by_origin_name = self._find_node_by_origin_name(input_node.name)
+            assert node_by_origin_name is not None, \
+                f"Prediction node with name {input_node.name} not found in model graph"
+            input_node_id = node_by_origin_name.id
         else:
             input_node_id = self._find_layer_node_id(input_node.name)
 

{leap_model_parser-0.1.185.dev12 → leap_model_parser-0.1.187}/leap_model_parser/model_parser.py

@@ -5,10 +5,10 @@ import tarfile
 import tempfile
 from importlib.util import find_spec
 from pathlib import Path
-from typing import Callable, Optional, List, Dict, Tuple, Type
+
 
 import tensorflow as tf  # type: ignore
-from code_loader.contract.mapping import NodeConnection, leap_output
+from code_loader.contract.mapping import NodeConnection  # type: ignore
 from keras import Model  # type: ignore
 from keras_data_format_converter import convert_channels_first_to_last  # type: ignore
 from leap_model_rebuilder import rebuild_model  # type: ignore
@@ -21,7 +21,9 @@ from tensorflow.keras.models import load_model  # type: ignore
 from leap_model_parser.contract.graph import Node, InputInfo
 from leap_model_parser.contract.importmodelresponse import ImportModelTypeEnum
 from leap_model_parser.keras_json_model_import import KerasJsonModelImport
-from leap_model_parser.leap_graph_editor import LeapGraphEditor
+from leap_model_parser.leap_graph_editor import LeapGraphEditor  # type: ignore
+
+from typing import Callable, Optional, List, Dict, Tuple, Type
 
 onnx_imported = False
 package_name = 'onnx'
@@ -33,7 +35,7 @@ if spec is not None:
 
 class ModelParser:
     def __init__(self, should_transform_inputs_and_outputs=False,
-                 custom_layers: Optional[Dict[str, Type[tf.keras.layers.Layer]]] = None,
+                 custom_layers=None,
                  mapping_connections: Optional[List[NodeConnection]] = None):
         self._should_transform_inputs_and_outputs = should_transform_inputs_and_outputs
         self.custom_layers = custom_layers

{leap_model_parser-0.1.185.dev12 → leap_model_parser-0.1.187}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "leap-model-parser"
-version = "0.1.185.dev12"
+version = "0.1.187"
 description = ""
 authors = ["idan <idan.yogev@tensorleap.ai>"]
 license = "MIT"
@@ -22,7 +22,7 @@ onnx2kerastl = "0.0.174"
 keras-data-format-converter = "0.1.22"
 leap-model-rebuilder = "0.1.7"
 tensorflow-io-gcs-filesystem = "0.34.0"
-code-loader = "1.0.87.dev8"
+code-loader = "1.0.89"
 
 [tool.poetry.dev-dependencies]
 pytest = "^7.1.1"