leap-model-parser 0.1.250.dev3__tar.gz → 0.1.250.dev5__tar.gz

{leap_model_parser-0.1.250.dev3 → leap_model_parser-0.1.250.dev5}/PKG-INFO

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

{leap_model_parser-0.1.250.dev3 → leap_model_parser-0.1.250.dev5}/leap_model_parser/leap_graph_editor.py

@@ -2,14 +2,13 @@ from enum import Enum
 from typing import Optional, Dict, Any, List
 
 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 LeapGraphEditor:
-    def __init__(self, model_graph: Dict[str, Node], keras_model: Model):
+    def __init__(self, model_graph: Dict[str, Node], keras_model=None):
         self.model_graph = model_graph
         self.keras_model = keras_model
 
@@ -80,6 +79,14 @@ class LeapGraphEditor:
                 return node
         return None
 
+    def _resolve_prediction_origin_name(self, prediction_index: int) -> str:
+        from keras import Model  # deferred — only Keras path calls this
+        return self.keras_model.outputs[prediction_index].node.layer.name
+
+    def _resolve_input_origin_name(self, input_index: int) -> str:
+        from keras import Model  # deferred — only Keras path calls this
+        return self.keras_model.inputs[input_index].node.layer.name
+
     def _replace_prediction_node_name_with_correct_name(self, connections: List[NodeConnection]) -> List[NodeConnection]:
         for connection in connections:
             if connection.node_inputs is None:
@@ -87,7 +94,7 @@ class LeapGraphEditor:
             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', ''))
-                    origin_name = self.keras_model.outputs[prediction_index].node.layer.name
+                    origin_name = self._resolve_prediction_origin_name(prediction_index)
                     input_node.name = origin_name
 
         return connections
@@ -227,7 +234,7 @@ class LeapGraphEditor:
 
     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
+        origin_name = self._resolve_input_origin_name(input_index)
         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

leap_model_parser-0.1.250.dev5/leap_model_parser/torch_graph_editor.py

@@ -0,0 +1,93 @@
+from typing import Dict, List, Optional
+
+from torch.export import ExportedProgram
+
+from leap_model_parser.contract.graph import Node
+from leap_model_parser.leap_graph_editor import LeapGraphEditor
+
+
+class TorchLeapGraphEditor(LeapGraphEditor):
+    """
+    PT2 variant of LeapGraphEditor. Resolves Prediction{N} indices from
+    ExportedProgram output specs instead of keras_model.outputs.
+    """
+
+    def __init__(self, model_graph: Dict[str, Node], ep: ExportedProgram):
+        self.model_graph = model_graph
+        self.keras_model = None
+        self._next_node_id_index = 0
+        self.ep = ep
+
+    def _resolve_prediction_origin_name(self, prediction_index: int) -> str:
+        module_path = self._find_pt2_output_module_path(prediction_index)
+        if module_path is None:
+            n = self._count_pt2_outputs()
+            raise ValueError(
+                f"Cannot resolve PT2 output module for Prediction{prediction_index}. "
+                f"ExportedProgram has {n} outputs."
+            )
+        return module_path
+
+    def _handle_input_node_with_index(self, input_node) -> str:
+        return self._add_input_encoder_not_connected_to_the_model_node(input_node.name)
+
+    def _resolve_input_origin_name(self, input_index: int) -> str:
+        from torch.export.graph_signature import InputKind
+        user_inputs = []
+        node_to_spec = {s.arg.name: s for s in self.ep.graph_signature.input_specs}
+        for node in self.ep.graph_module.graph.nodes:
+            if node.op != "placeholder":
+                break
+            spec = node_to_spec.get(node.name)
+            if spec and spec.kind not in (InputKind.PARAMETER, InputKind.BUFFER):
+                user_inputs.append(node.name)
+        if input_index >= len(user_inputs):
+            raise ValueError(
+                f"Cannot resolve PT2 input at index {input_index}. "
+                f"ExportedProgram has {len(user_inputs)} user inputs."
+            )
+        return user_inputs[input_index]
+
+    def _find_pt2_output_module_path(self, output_index: int) -> Optional[str]:
+        output_node = None
+        for node in self.ep.graph_module.graph.nodes:
+            if node.op == "output":
+                output_node = node
+                break
+        if output_node is None:
+            return None
+
+        args = output_node.args[0]
+        if not isinstance(args, (list, tuple)):
+            args = [args]
+
+        if output_index >= len(args):
+            return None
+
+        output_ref = args[output_index]
+        if output_ref is None or not hasattr(output_ref, "meta"):
+            return None
+
+        visited = set()
+        current = output_ref
+        while current.op == "call_function" and current not in visited:
+            stack = current.meta.get("nn_module_stack", {})
+            if stack:
+                return list(stack.items())[-1][1][0]
+            if current.args:
+                visited.add(current)
+                current = current.args[0]
+            else:
+                break
+
+        stack = current.meta.get("nn_module_stack", {})
+        if stack:
+            return list(stack.items())[-1][1][0]
+        return None
+
+    def _count_pt2_outputs(self) -> int:
+        for node in self.ep.graph_module.graph.nodes:
+            if node.op == "output":
+                args = node.args[0]
+                return len(args) if isinstance(args, (list, tuple)) else 1
+        return 0

{leap_model_parser-0.1.250.dev3 → leap_model_parser-0.1.250.dev5}/leap_model_parser/torch_model_parser.py

@@ -1,4 +1,4 @@
-from typing import Dict, List, Tuple
+from typing import Dict, List, Optional, Tuple
 
 import torch
 from torch.export import ExportedProgram, load
@@ -23,11 +23,26 @@ class TorchModelParser:
                 id=mod_path,
                 name=mod_path.split(".")[-1] if "." in mod_path else mod_path,
                 position=[i * 100, 0],
-                data={"type": class_name, "framework": "torch"},
+                data={
+                    "type": "Layer",
+                    "class_name": class_name,
+                    "framework": "torch",
+                    "origin_name": mod_path,
+                },
             )
             nodes[mod_path] = node
         return nodes
 
+    def apply_mapping_connections(
+        self,
+        ui_nodes: Dict[str, Node],
+        ep: ExportedProgram,
+        mapping_connections: List,
+    ) -> None:
+        from leap_model_parser.torch_graph_editor import TorchLeapGraphEditor
+        editor = TorchLeapGraphEditor(ui_nodes, ep)
+        editor.add_connections_to_graph(mapping_connections)
+
     def _extract_input_infos(self, ep: ExportedProgram) -> List[InputInfo]:
         node_to_spec = {s.arg.name: s for s in ep.graph_signature.input_specs}
         infos: List[InputInfo] = []

{leap_model_parser-0.1.250.dev3 → leap_model_parser-0.1.250.dev5}/pyproject.toml

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