leap-model-parser 0.1.185.dev4__tar.gz → 0.1.185.dev6__tar.gz

{leap_model_parser-0.1.185.dev4 → leap_model_parser-0.1.185.dev6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: leap-model-parser
-Version: 0.1.185.dev4
+Version: 0.1.185.dev6
 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.dev4)
+Requires-Dist: code-loader (==1.0.87.dev5)
 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.dev4 → leap_model_parser-0.1.185.dev6}/leap_model_parser/leap_graph_editor.py

@@ -102,7 +102,8 @@ class LeapGraphEditor:
             #     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.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:
@@ -231,7 +232,7 @@ class LeapGraphEditor:
         self.model_graph[new_node_id] = metric_node
         return new_node_id
 
-    def _add_loss_node(self, loss_name: str, is_custom_loss: bool) -> str:
+    def _add_loss_node(self, loss_name: str, is_custom_loss: bool, arg_names: Optional[List[str]]=None) -> str:
         new_node_id = self._generate_new_node_id()
 
         loss_type = 'CustomLoss' if is_custom_loss else 'Loss'
@@ -250,6 +251,9 @@ class LeapGraphEditor:
             #     f'{new_node_id}-loss': {'connections': []}
             # }
         )
+        if arg_names is not None:
+            loss_node.data['arg_names'] = arg_names
+
 
         self.model_graph[new_node_id] = loss_node
         return new_node_id

{leap_model_parser-0.1.185.dev4 → leap_model_parser-0.1.185.dev6}/leap_model_parser/model_parser.py

@@ -8,7 +8,7 @@ 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
+from code_loader.contract.mapping import NodeConnection, leap_output
 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
@@ -31,7 +31,6 @@ if spec is not None:
 
     onnx_imported = True
 
-
 class ModelParser:
     def __init__(self, should_transform_inputs_and_outputs=False,
                  custom_layers: Optional[Dict[str, Type[tf.keras.layers.Layer]]] = None,

{leap_model_parser-0.1.185.dev4 → leap_model_parser-0.1.185.dev6}/pyproject.toml

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