leap-model-parser 0.1.248.dev0__tar.gz → 0.1.250.dev0__tar.gz

{leap_model_parser-0.1.248.dev0 → leap_model_parser-0.1.250.dev0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: leap-model-parser
-Version: 0.1.248.dev0
+Version: 0.1.250.dev0
 Summary: 
 Home-page: https://github.com/tensorleap/leap-model-parser
 License: MIT
@@ -17,10 +17,11 @@ 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.190.dev0)
+Requires-Dist: onnx2kerastl (==0.0.190)
 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.248.dev0 → leap_model_parser-0.1.250.dev0}/leap_model_parser/contract/graph.py

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

{leap_model_parser-0.1.248.dev0 → leap_model_parser-0.1.250.dev0}/leap_model_parser/contract/importmodelresponse.py

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

leap_model_parser-0.1.250.dev0/leap_model_parser/torch_model_parser.py

@@ -0,0 +1,42 @@
+from typing import Dict, List, 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": class_name, "framework": "torch"},
+            )
+            nodes[mod_path] = node
+        return nodes
+
+    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.dev0/leap_model_parser/torch_utils.py

@@ -0,0 +1,65 @@
+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, 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]]:
+    mapping: Dict[str, Tuple[str, str]] = {}
+    for node in ep.graph_module.graph.nodes:
+        if node.op != "call_function":
+            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):
+        super().__init__(gm)
+        self.cache: Dict[str, torch.Tensor] = {}
+
+    def run_node(self, node: fx.Node) -> Any:
+        result = super().run_node(node)
+        if isinstance(result, torch.Tensor):
+            self.cache[node.name] = result.detach().clone()
+        return result

{leap_model_parser-0.1.248.dev0 → leap_model_parser-0.1.250.dev0}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "leap-model-parser"
-version = "0.1.248.dev0"
+version = "0.1.250.dev"
 description = ""
 authors = ["idan <idan.yogev@tensorleap.ai>"]
 license = "MIT"
@@ -18,11 +18,12 @@ 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.190.dev0"
+onnx2kerastl = "0.0.190"
 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"