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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: leap-model-parser
-Version: 0.1.185.dev11
+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.dev7)
+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.dev11 → leap_model_parser-0.1.187}/leap_model_parser/leap_graph_editor.py

@@ -1,51 +1,12 @@
 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
 
 
-# class NodeMappingType(Enum):
-#     Visualizer = 'Visualizer'
-#     Metric = 'Metric'
-#     GroundTruth = 'GroundTruth'
-#     Input = 'Input'
-#     Layer = 'Layer'
-#     Loss = 'Loss'
-#     CustomLoss = 'CustomLoss'
-#     Optimizer = 'Optimizer'
-#     Prediction0 = 'Prediction0'
-#     Prediction1 = 'Prediction1'
-#     Prediction2 = 'Prediction2'
-#     Prediction3 = 'Prediction3'
-#     Input0 = 'Input0'
-#     Input1 = 'Input1'
-#     Input2 = 'Input2'
-#     Input3 = 'Input3'
-#     Input4 = 'Input4'
-#     Input5 = 'Input5'
-#
-#
-#
-#
-#
-# @dataclass
-# class NodeMapping:
-#     name: str
-#     type: NodeMappingType
-#     user_unique_name: Optional[str] = None
-#     sub_type: Optional[str] = None
-#     arg_names: Optional[List[str]] = None
-#
-#
-# @dataclass
-# class NodeConnection:
-#     node: NodeMapping
-#     node_inputs: Optional[Dict[str, NodeMapping]]
-#     prediction_type_name: Optional[str] = None
-
 
 class LeapGraphEditor:
     def __init__(self, model_graph: Dict[str, Node], keras_model: Model):
@@ -55,33 +16,21 @@ class LeapGraphEditor:
         node_ids_as_int = [int(node_id) for node_id in model_graph.keys()]
         self._next_node_id_index = max(node_ids_as_int) + 1
 
-    # def add_dataset(self, dataset_name: str, raw_dataset_version: Dict[str, Any],
-    #                 dataset_parse_result: DatasetIntegParseResult):
-    #
-    #     LeapGraphEditor._add_setup_to_metadata(raw_dataset_version['metadata'], dataset_parse_result)
-    #     raw_dataset_version['name'] = dataset_name
-    #
-    #     dataset_node = self._get_dataset_node()
-    #     dataset_node.data['datasetVersion'] = raw_dataset_version
-    #     dataset_node.data['selected_dataset'] = dataset_name
-    #     self._add_arg_names_to_visualizers(dataset_parse_result)
 
     def add_connections_to_graph(self, connections: List[NodeConnection]):
-        connections = self._validate_and_reorder_connections_list(connections)
+        connections = self._replace_prediction_node_name_with_correct_name(connections)
         for connection in connections:
             self._add_node_connection_to_graph(connection)
 
     def _add_node_connection_to_graph(self, node_connection: NodeConnection):
-        # if node_connection.node.type.value.startswith('Input'):
-        #     input_index = int(node_connection.node.type.value.replace('Input', ''))
-        #
-        #     origin_name = self.keras_model.inputs[input_index].node.layer.name
-        #
-        #     _find_node_by_origin_name
-        # # elif node_connection.node.type == NodeMappingType.Input:
-        # #     self._find_or_add_input_node()
-
-
+        if node_connection.node.type == NodeMappingType.PredictionLabels:
+            prediction_labels_name = node_connection.node.name
+            prediction_node = list(node_connection.node_inputs.values())[0]
+            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(
                 node_connection.node.name, node_connection.node.sub_type,
@@ -97,26 +46,12 @@ class LeapGraphEditor:
                 input_node_id = self._find_or_add_input_node(node)
                 self._add_connection_to_node(new_metric_node_id, input_name, input_node_id)
         elif node_connection.node.type in (NodeMappingType.Loss, NodeMappingType.CustomLoss):
-            prediction_type_name = node_connection.prediction_type_name
-            # if prediction_type_name is None:
-            #     raise Exception("prediction_type_name is required for loss connection")
-
             new_loss_node_id = self._add_loss_node(node_connection.node.name,
                                                    node_connection.node.type == NodeMappingType.CustomLoss,
                                                    node_connection.node.arg_names)
             for input_name, node in node_connection.node_inputs.items():
                 input_node_id = self._find_or_add_input_node(node)
-                # if node.type == NodeMappingType.Layer:
-                #     self.model_graph[input_node_id].data['prediction_type'] = prediction_type_name
                 self._add_connection_to_node(new_loss_node_id, input_name, input_node_id)
-        # elif node_connection.node.type == NodeMappingType.Optimizer:
-        #     new_optimizer_node_id = self._add_optimizer_node(node_connection.node.name)
-        #     loss_node_ids = self._get_all_loss_node_ids()
-        #     assert len(loss_node_ids) > 0
-        #     for i, loss_node_id in enumerate(loss_node_ids):
-        #         self._add_connection_to_node(new_optimizer_node_id, str(i), loss_node_id)
-        #     self.model_graph[new_optimizer_node_id].data['custom_input_keys'] = list(
-        #         self.model_graph[new_optimizer_node_id].inputs.keys())
         else:
             raise Exception(f"Can't add node of type {node_connection.node.type.name}")
 
@@ -140,30 +75,18 @@ class LeapGraphEditor:
                 return node
         return None
 
-    def _validate_and_reorder_connections_list(self, connections: List[NodeConnection]) -> List[NodeConnection]:
-        # optimizers = [connection for connection in connections if connection.node.type == NodeType.Optimizer]
+    def _replace_prediction_node_name_with_correct_name(self, connections: List[NodeConnection]) -> List[NodeConnection]:
         for connection in connections:
             if connection.node_inputs is None:
                 continue
             for input_name, input_node in connection.node_inputs.items():
                 if 'Prediction' in input_node.type.value:
-                    prediction_index= int(input_node.type.value.replace('Prediction', ''))
+                    prediction_index = int(input_node.type.value.replace('Prediction', ''))
                     origin_name = self.keras_model.outputs[prediction_index].node.layer.name
                     input_node.name = origin_name
 
         return connections
-        losses = [connection for connection in connections
-                  if connection.node.type in (NodeMappingType.Loss, NodeMappingType.CustomLoss)]
-        visualizers = [connection for connection in connections if connection.node.type == NodeMappingType.Visualizer]
 
-        # if len(optimizers) == 0:
-        #     raise Exception('At least one optimizer needed')
-        # if len(losses) == 0:
-        #     raise Exception('At least one loss needed')
-        # if len(optimizers) + len(losses) + len(visualizers) < len(connections):
-        #     raise Exception('Unsupported node type')
-
-        return visualizers + losses
 
     def _find_encoder_node_id(self, encoder_name: str) -> Optional[str]:
         for node_id, node_response in self.model_graph.items():
@@ -247,9 +170,6 @@ class LeapGraphEditor:
             outputs={
                 f'{new_node_id}-loss': ConnectionOutput([])
             }
-            # outputs={
-            #     f'{new_node_id}-loss': {'connections': []}
-            # }
         )
         if arg_names is not None:
             loss_node.data['arg_names'] = arg_names
@@ -258,28 +178,12 @@ class LeapGraphEditor:
         self.model_graph[new_node_id] = loss_node
         return new_node_id
 
-    # def _add_optimizer_node(self, optimizer_name: str) -> str:
-    #     new_node_id = self._generate_new_node_id()
-    #
-    #     optimizer_node = NodeResponse(
-    #         new_node_id,
-    #         optimizer_name,
-    #         data={'type': 'Optimizer', 'selected': optimizer_name},
-    #         inputs={},
-    #         outputs={})
-    #
-    #     self.model_graph[new_node_id] = optimizer_node
-    #     return new_node_id
-
     def _get_output_name_from_node_id(self, input_node_id: str, input_name: Optional[str] = None) -> str:
         input_node_outputs_len = len(self.model_graph[input_node_id].outputs)
         if input_node_outputs_len == 0:
             output_name_to_add = f'{input_node_id}-feature_map'
             self.model_graph[input_node_id].outputs[output_name_to_add] = ConnectionOutput([])
 
-            # self.model_graph[input_node_id].outputs[output_name_to_add] = {
-            #     'connections': []
-            # }
             return output_name_to_add
         if input_node_outputs_len == 1:
             return list(self.model_graph[input_node_id].outputs.keys())[0]
@@ -296,16 +200,32 @@ class LeapGraphEditor:
         output_name = self._get_output_name_from_node_id(input_node_id, input_name)
         input_name = f'{node_id}-{input_name}'
         self.model_graph[node_id].inputs[input_name] = ConnectionInput([InputData(input_node_id, output_name)])
-        # self.model_graph[node_id].inputs[input_name] = {
-        #     'connections': [{'data': {}, 'node': input_node_id, 'output': output_name}]
-        # }
 
-        # if 'connections' not in self.model_graph[input_node_id].outputs[output_name]:
-        #     self.model_graph[input_node_id].outputs[output_name]['connections'] = []
         output_connection = OutputData(node_id, input_name)
-        # output_connection = {'input': input_name, 'node': node_id, 'data': {}}
         self.model_graph[input_node_id].outputs[output_name].connections.append(output_connection)
 
+    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 = 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:
+            new_output_key = f'{input_node_id}-{input_node.name}'
+            if output_key == new_output_key:
+                continue
+            self.model_graph[input_node_id].outputs[new_output_key] = self.model_graph[input_node_id].outputs[
+                output_key]
+            del self.model_graph[input_node_id].outputs[output_key]
+            for connection in self.model_graph[input_node_id].outputs[new_output_key].connections:
+                for connection_input in self.model_graph[connection.node].inputs[connection.input].connections:
+                    if connection_input.output == output_key:
+                        connection_input.output = new_output_key
+
+        return input_node_id
+
     def _find_or_add_input_node(self, input_node: NodeMapping) -> str:
         if input_node.type in (NodeMappingType.Input, NodeMappingType.GroundTruth):
             input_node_id = self._find_encoder_node_id(input_node.name)
@@ -315,26 +235,12 @@ class LeapGraphEditor:
                 else:
                     raise Exception(f'Couldnt find input node name {input_node.name}')
         elif 'Input' in input_node.type.value:
-            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
-            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}'
-                if output_key == new_output_key:
-                    continue
-                self.model_graph[input_node_id].outputs[new_output_key] = self.model_graph[input_node_id].outputs[output_key]
-                del self.model_graph[input_node_id].outputs[output_key]
-                for connection in self.model_graph[input_node_id].outputs[new_output_key].connections:
-                    for connection_input in self.model_graph[connection.node].inputs[connection.input].connections:
-                        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}")
+            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)
 
@@ -350,13 +256,6 @@ class LeapGraphEditor:
                 loss_node_ids.append(node_id)
         return loss_node_ids
 
-    # def _get_dataset_node(self) -> NodeResponse:
-    #     for node_response in self.model_graph.values():
-    #         if 'type' in node_response.data and node_response.data['type'] == 'dataset':
-    #             return node_response
-    #
-    #     raise Exception("Didn't find dataset node")
-
     @staticmethod
     def _convert_dataclass_to_json_dict(_dataclass):
         if isinstance(_dataclass, Enum):
@@ -372,20 +271,3 @@ class LeapGraphEditor:
                 for element in _dataclass
             ]
         return _dataclass
-
-    # @staticmethod
-    # def _add_setup_to_metadata(dataset_version_metadata: Dict[str, Any],
-    #                            dataset_parse_result: DatasetIntegParseResult):
-    #     setup_json = LeapGraphEditor._convert_dataclass_to_json_dict(dataset_parse_result.setup)
-    #
-    #     dataset_version_metadata['setup'] = setup_json
-
-    # def _add_arg_names_to_visualizers(self, dataset_parse_result: DatasetIntegParseResult):
-    #     visualizer_instance_by_name: Dict[str, VisualizerInstance] = {
-    #         visualizer_instance.name: visualizer_instance
-    #         for visualizer_instance in dataset_parse_result.setup.visualizers
-    #     }
-    #
-    #     for _, node_response in self.model_graph.items():
-    #         if node_response.data['type'] == 'Visualizer':
-    #             node_response.data['arg_names'] = visualizer_instance_by_name[node_response.data['selected']].arg_names

{leap_model_parser-0.1.185.dev11 → 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.dev11 → leap_model_parser-0.1.187}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "leap-model-parser"
-version = "0.1.185.dev11"
+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.dev7"
+code-loader = "1.0.89"
 
 [tool.poetry.dev-dependencies]
 pytest = "^7.1.1"