leap-model-parser 0.1.197__tar.gz → 0.1.198.dev1__tar.gz

{leap_model_parser-0.1.197 → leap_model_parser-0.1.198.dev1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: leap-model-parser
-Version: 0.1.197
+Version: 0.1.198.dev1
 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.99)
+Requires-Dist: code-loader (==1.0.100.dev1)
 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.197 → leap_model_parser-0.1.198.dev1}/leap_model_parser/leap_graph_editor.py

@@ -31,7 +31,7 @@ class LeapGraphEditor:
                 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:
+        elif node_connection.node.type == NodeMappingType.Visualizer:
             new_visualizer_node_id = self._add_visualizer_node(
                 node_connection.node.name, node_connection.node.sub_type,
                 node_connection.node.user_unique_name, node_connection.node.arg_names)
@@ -71,6 +71,8 @@ class LeapGraphEditor:
 
     def _find_input_node_by_origin_name(self, origin_name: str) -> Optional[Node]:
         for node in self.model_graph.values():
+            if node.data.get('original_output_name') == origin_name:
+                return node
             if node.data.get('output_name') == origin_name:
                 return node
         return None
@@ -210,7 +212,10 @@ class LeapGraphEditor:
         input_node_by_origin = self._find_input_node_by_origin_name(origin_name)
         assert input_node_by_origin is not None, f"Input node with origin name {origin_name} not found in model graph"
         input_node_id = input_node_by_origin.id
+        if 'original_output_name' not in self.model_graph[input_node_id].data:
+            self.model_graph[input_node_id].data['original_output_name'] = self.model_graph[input_node_id].data['output_name']
         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:
             new_output_key = f'{input_node_id}-{input_node.name}'

{leap_model_parser-0.1.197 → leap_model_parser-0.1.198.dev1}/leap_model_parser/model_parser.py

@@ -1,6 +1,7 @@
 import glob
 import json
 import ntpath
+import pickle
 import tarfile
 import tempfile
 from importlib.util import find_spec
@@ -8,7 +9,7 @@ from pathlib import Path
 
 
 import tensorflow as tf  # type: ignore
-from code_loader.contract.mapping import NodeConnection  # type: ignore
+from code_loader.contract.mapping import NodeConnection, NodeMapping  # 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

{leap_model_parser-0.1.197 → leap_model_parser-0.1.198.dev1}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "leap-model-parser"
-version = "0.1.197"
+version = "0.1.198.dev1"
 description = ""
 authors = ["idan <idan.yogev@tensorleap.ai>"]
 license = "MIT"
@@ -22,7 +22,7 @@ onnx2kerastl = "0.0.175"
 keras-data-format-converter = "0.1.22"
 leap-model-rebuilder = "0.1.7"
 tensorflow-io-gcs-filesystem = "0.34.0"
-code-loader = "1.0.99"
+code-loader = "1.0.100.dev1"
 
 [tool.poetry.dev-dependencies]
 pytest = "^7.1.1"