leap-model-parser 0.1.250.dev8__tar.gz → 0.1.250.dev10__tar.gz

{leap_model_parser-0.1.250.dev8 → leap_model_parser-0.1.250.dev10}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: leap-model-parser
-Version: 0.1.250.dev8
+Version: 0.1.250.dev10
 Summary: 
 Home-page: https://github.com/tensorleap/leap-model-parser
 License: MIT
@@ -17,11 +17,10 @@ Requires-Dist: keras-data-format-converter (==0.1.24)
 Requires-Dist: leap-model-rebuilder (==0.1.7)
 Requires-Dist: numpy (>=1.22.3,<2.0.0)
 Requires-Dist: onnx (==1.13.0)
-Requires-Dist: onnx2kerastl (==0.0.191)
+Requires-Dist: onnx2kerastl (==0.0.192.dev1)
 Requires-Dist: tensorflow (==2.12.0) ; platform_machine == "x86_64"
 Requires-Dist: tensorflow-io-gcs-filesystem (==0.34.0)
 Requires-Dist: tensorflow-macos (==2.12.0) ; platform_machine == "arm64"
-Requires-Dist: torch (>=2.3.0,<2.4.0)
 Project-URL: Repository, https://github.com/tensorleap/leap-model-parser
 Description-Content-Type: text/markdown
 

{leap_model_parser-0.1.250.dev8 → leap_model_parser-0.1.250.dev10}/leap_model_parser/contract/graph.py

@@ -63,5 +63,4 @@ class Node:
 @dataclass
 class ModelGraph:
     id: str
-    nodes: Dict[str, Node] = field(default_factory=dict)
-    framework: str = "keras"
+    nodes: Dict[str, Node] = field(default_factory=dict)

{leap_model_parser-0.1.250.dev8 → leap_model_parser-0.1.250.dev10}/leap_model_parser/contract/importmodelresponse.py

@@ -6,4 +6,3 @@ class ImportModelTypeEnum(Enum):
     ONNX = "ONNX"
     PB_TF2 = "PB_TF2"
     H5_TF2 = "H5_TF2"
-    PYTORCH_PT2 = "PYTORCH_PT2"

{leap_model_parser-0.1.250.dev8 → leap_model_parser-0.1.250.dev10}/leap_model_parser/leap_graph_editor.py

@@ -2,13 +2,14 @@ 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=None):
+    def __init__(self, model_graph: Dict[str, Node], keras_model: Model):
         self.model_graph = model_graph
         self.keras_model = keras_model
 
@@ -79,14 +80,6 @@ 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:
@@ -94,7 +87,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._resolve_prediction_origin_name(prediction_index)
+                    origin_name = self.keras_model.outputs[prediction_index].node.layer.name
                     input_node.name = origin_name
 
         return connections
@@ -234,7 +227,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._resolve_input_origin_name(input_index)
+        origin_name = self.keras_model.inputs[input_index].node.layer.name
         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.dev8 → leap_model_parser-0.1.250.dev10}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "leap-model-parser"
-version = "0.1.250.dev8"
+version = "0.1.250.dev10"
 description = ""
 authors = ["idan <idan.yogev@tensorleap.ai>"]
 license = "MIT"
@@ -18,12 +18,11 @@ tensorflow = {version = "2.12.0", markers = "platform_machine  == 'x86_64'"}
 tensorflow-macos = {version = "2.12.0", markers = "platform_machine  == 'arm64'"}
 numpy = "^1.22.3"
 onnx = "1.13.0"
-onnx2kerastl = "0.0.191"
+onnx2kerastl = "0.0.192.dev1"
 keras-data-format-converter = "0.1.24"
 leap-model-rebuilder = "0.1.7"
 tensorflow-io-gcs-filesystem = "0.34.0"
 code-loader = ">=1.0.127"
-torch = "~=2.3.0"
 
 [tool.poetry.dev-dependencies]
 pytest = "^7.1.1"

leap_model_parser-0.1.250.dev8/leap_model_parser/torch_graph_editor.py

@@ -1,93 +0,0 @@
-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.dev8/leap_model_parser/torch_model_parser.py

@@ -1,57 +0,0 @@
-from typing import Dict, List, Optional, Tuple
-
-import torch
-from torch.export import ExportedProgram, load
-from torch.export.graph_signature import InputKind
-
-from leap_model_parser.contract.graph import InputInfo, Node
-from leap_model_parser.torch_utils import build_module_to_node_map
-
-
-class TorchModelParser:
-    def parse(self, file_path: str) -> Tuple[Dict[str, Node], List[InputInfo], ExportedProgram, str]:
-        ep = load(file_path)
-        nodes = self._build_node_graph(ep)
-        input_infos = self._extract_input_infos(ep)
-        return nodes, input_infos, ep, "torch"
-
-    def _build_node_graph(self, ep: ExportedProgram) -> Dict[str, Node]:
-        m2n = build_module_to_node_map(ep)
-        nodes: Dict[str, Node] = {}
-        for i, (mod_path, (node_name, class_name)) in enumerate(m2n.items()):
-            node = Node(
-                id=mod_path,
-                name=mod_path.split(".")[-1] if "." in mod_path else mod_path,
-                position=[i * 100, 0],
-                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] = []
-        for node in 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):
-                val = node.meta.get("val")
-                shape = list(val.shape) if val is not None else []
-                infos.append(InputInfo(name=node.name, shape=shape))
-        return infos

leap_model_parser-0.1.250.dev8/leap_model_parser/torch_utils.py

@@ -1,149 +0,0 @@
-import numpy as np
-import torch
-import torch.fx as fx
-from torch.export import ExportedProgram
-from torch.export.graph_signature import InputKind
-from typing import Any, Dict, List, Optional, Tuple
-
-
-def get_full_args(ep: ExportedProgram, user_inputs: List[torch.Tensor]) -> List[torch.Tensor]:
-    node_to_spec = {s.arg.name: s for s in ep.graph_signature.input_specs}
-    user_iter = iter(user_inputs)
-    constants = getattr(ep, "constants", {}) or {}
-    full_args = []
-    for node in ep.graph_module.graph.nodes:
-        if node.op != "placeholder":
-            break
-        spec = node_to_spec.get(node.name)
-        if spec and spec.kind in (InputKind.PARAMETER, InputKind.BUFFER):
-            if spec.target in ep.state_dict:
-                full_args.append(ep.state_dict[spec.target].clone())
-            else:
-                full_args.append(constants[spec.target].clone())
-        else:
-            full_args.append(next(user_iter))
-    return full_args
-
-
-def build_module_to_node_map(ep: ExportedProgram) -> Dict[str, Tuple[str, str]]:
-    import operator as _op
-    mapping: Dict[str, Tuple[str, str]] = {}
-    for node in ep.graph_module.graph.nodes:
-        if node.op != "call_function":
-            continue
-        if node.target is _op.getitem:
-            continue
-        stack = node.meta.get("nn_module_stack", {})
-        if not stack:
-            continue
-        for mod_path, mod_type in stack.values():
-            class_name = mod_type.__name__ if not isinstance(mod_type, str) else mod_type.split(".")[-1]
-            mapping[mod_path] = (node.name, class_name)
-    return mapping
-
-
-def extract_tensor(output: Any) -> torch.Tensor:
-    if isinstance(output, torch.Tensor):
-        return output
-    if isinstance(output, (tuple, list)):
-        t = output[0]
-        return t if isinstance(t, torch.Tensor) else extract_tensor(t)
-    if hasattr(output, "pooler_output") and output.pooler_output is not None:
-        return output.pooler_output
-    if hasattr(output, "last_hidden_state"):
-        return output.last_hidden_state[:, 0, :]
-    if hasattr(output, "logits"):
-        return output.logits
-    return output[0]
-
-
-class ActivationInterpreter(fx.Interpreter):
-    def __init__(self, gm: fx.GraphModule, grad_mode: bool = False):
-        super().__init__(gm)
-        self.cache: Dict[str, torch.Tensor] = {}
-        self._grad_mode = grad_mode
-
-    def _store_tensor(self, name: str, value: Any) -> None:
-        if isinstance(value, torch.Tensor):
-            if self._grad_mode:
-                self.cache[name] = value
-                if value.requires_grad:
-                    value.retain_grad()
-            else:
-                self.cache[name] = value.detach().clone()
-        elif isinstance(value, (tuple, list)):
-            for item in value:
-                if isinstance(item, torch.Tensor):
-                    if self._grad_mode:
-                        self.cache[name] = item
-                        if item.requires_grad:
-                            item.retain_grad()
-                    else:
-                        self.cache[name] = item.detach().clone()
-                    break
-
-    def run_node(self, node: fx.Node) -> Any:
-        result = super().run_node(node)
-        self._store_tensor(node.name, result)
-        return result
-
-
-def get_pt2_output_module_paths(ep: ExportedProgram) -> List[Optional[str]]:
-    """Return the nn.Module path responsible for each output slot in the FX graph."""
-    output_node = None
-    for node in ep.graph_module.graph.nodes:
-        if node.op == "output":
-            output_node = node
-            break
-    if output_node is None:
-        return []
-
-    args = output_node.args[0]
-    if not isinstance(args, (list, tuple)):
-        args = [args]
-
-    result: List[Optional[str]] = []
-    for output_ref in args:
-        if output_ref is None or not hasattr(output_ref, "meta"):
-            result.append(None)
-            continue
-
-        visited: set = set()
-        current = output_ref
-        mod_path = None
-        while current.op == "call_function" and current not in visited:
-            stack = current.meta.get("nn_module_stack", {})
-            if stack:
-                mod_path = list(stack.items())[-1][1][0]
-                break
-            if current.args:
-                visited.add(current)
-                current = current.args[0]
-            else:
-                break
-
-        if mod_path is None:
-            stack = current.meta.get("nn_module_stack", {})
-            if stack:
-                mod_path = list(stack.items())[-1][1][0]
-
-        result.append(mod_path)
-
-    return result
-
-
-def get_module_activations(
-    ep: ExportedProgram,
-    user_inputs: List[torch.Tensor],
-) -> Dict[str, np.ndarray]:
-    tensors = [torch.from_numpy(x.astype(np.float32)) if isinstance(x, np.ndarray) else x
-               for x in user_inputs]
-    m2n = build_module_to_node_map(ep)
-    interp = ActivationInterpreter(ep.graph_module)
-    with torch.no_grad():
-        interp.run(*get_full_args(ep, tensors))
-    result: Dict[str, np.ndarray] = {}
-    for mod_path, (node_name, _) in m2n.items():
-        if node_name in interp.cache:
-            result[mod_path] = interp.cache[node_name].numpy()
-    return result