kinfer 0.0.3__tar.gz → 0.0.4__tar.gz

{kinfer-0.0.3/kinfer.egg-info → kinfer-0.0.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: kinfer
-Version: 0.0.3
+Version: 0.0.4
 Summary: The kinfer project
 Home-page: https://github.com/kscalelabs/kinfer.git
 Author: K-Scale Labs

{kinfer-0.0.3 → kinfer-0.0.4}/kinfer/__init__.py

@@ -1,3 +1,3 @@
-__version__ = "0.0.3"
+__version__ = "0.0.4"
 
 from . import export, inference

{kinfer-0.0.3 → kinfer-0.0.4}/kinfer/export/pytorch.py

@@ -1,6 +1,8 @@
 """PyTorch model export utilities."""
 
 import inspect
+import json
+import logging
 import sys
 from dataclasses import fields, is_dataclass
 from io import BytesIO
@@ -30,7 +32,6 @@ def get_model_info(model: nn.Module) -> Dict[str, Any]:
         if name == "self":
             continue
         params_info[name] = {
-            "annotation": str(param.annotation),
             "default": None if param.default is param.empty else str(param.default),
         }
 
@@ -53,19 +54,21 @@ def add_metadata_to_onnx(
     Returns:
         ONNX model with added metadata
     """
-    # Add model metadata
-    for key, value in metadata.items():
-        meta = model_proto.metadata_props.add()
-        meta.key = key
-        meta.value = str(value)
+    # Build metadata dictionary
+    metadata_dict = metadata.copy()
 
     # Add configuration if provided
-    if config is not None and is_dataclass(config):
-        for field in fields(config):
-            value = getattr(config, field.name)
-            meta = model_proto.metadata_props.add()
-            meta.key = field.name
-            meta.value = str(value)
+    if config is not None:
+        if is_dataclass(config):
+            for field in fields(config):
+                metadata_dict[field.name] = getattr(config, field.name)
+        elif not isinstance(config, dict):
+            raise ValueError("config must be a dataclass or dict. Got: " + str(type(config)))
+
+    # Add metadata as JSON string
+    meta = model_proto.metadata_props.add()
+    meta.key = "kinfer_metadata"
+    meta.value = json.dumps(metadata_dict)
 
     return model_proto
 
@@ -165,6 +168,10 @@ def export_to_onnx(
 
     # Create example inputs if not provided
     if input_tensors is None:
+        logging.warning(
+            "No input_tensors provided. Attempting to automatically infer input shapes. "
+            "Note: Input shape inference is *highly* experimental and may not work correctly for all models."
+        )
         try:
             input_tensors = create_example_inputs(model)
             model_info["inferred_input_shapes"] = str(
@@ -198,6 +205,10 @@ def export_to_onnx(
     # Load as ONNX model
     model_proto = onnx.load_model(buffer)
 
+    # Add config dict to model info if provided
+    if isinstance(config, dict):
+        model_info.update(config)
+
     # Add metadata
     model_proto = add_metadata_to_onnx(model_proto, model_info, config)
 

{kinfer-0.0.3 → kinfer-0.0.4}/kinfer/inference/python.py

@@ -1,5 +1,7 @@
 """ONNX model inference utilities for Python."""
 
+import json
+import logging
 from typing import Any, Dict, List, Union
 
 import numpy as np
@@ -28,8 +30,21 @@ class ONNXModel:
             model_path,
         )
 
-        # Extract metadata
-        self.metadata = {prop.key: prop.value for prop in self.model.metadata_props}
+        # Extract metadata and attempt to parse JSON values
+        self.metadata = {}
+        self.attached_metadata = {}
+        for prop in self.model.metadata_props:
+            if prop.key == "kinfer_metadata":
+                try:
+                    self.metadata = json.loads(prop.value)
+                except json.JSONDecodeError:
+                    logging.warning(
+                        "Failed to parse kinfer_metadata value with JSON parser. Saving as string: %s",
+                        prop.value,
+                    )
+                    self.metadata = prop.value
+
+            self.attached_metadata[prop.key] = prop.value
 
         # Get input and output details
         self.input_details = [{"name": x.name, "shape": x.shape, "type": x.type} for x in self.session.get_inputs()]
@@ -66,7 +81,7 @@ class ONNXModel:
         else:
             return {detail["name"]: arr for detail, arr in zip(self.output_details, outputs)}
 
-    def get_metadata(self) -> Dict[str, str]:
+    def get_metadata(self) -> Dict[str, Any]:
         """Get model metadata.
 
         Returns:

{kinfer-0.0.3 → kinfer-0.0.4/kinfer.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: kinfer
-Version: 0.0.3
+Version: 0.0.4
 Summary: The kinfer project
 Home-page: https://github.com/kscalelabs/kinfer.git
 Author: K-Scale Labs

{kinfer-0.0.3 → kinfer-0.0.4}/tests/test_infer.py

@@ -67,8 +67,8 @@ def test_model_metadata(model_path: str) -> None:
     # Check if config parameters are in metadata
     assert "hidden_size" in metadata
     assert "num_layers" in metadata
-    assert metadata["hidden_size"] == "64"
-    assert metadata["num_layers"] == "2"
+    assert metadata["hidden_size"] == 64
+    assert metadata["num_layers"] == 2
 
 
 def test_model_inference(model_path: str) -> None:
@@ -123,8 +123,8 @@ def test_comprehensive_model_workflow(tmp_path: Path) -> None:
     metadata = onnx_model.get_metadata()
     assert "hidden_size" in metadata
     assert "num_layers" in metadata
-    assert metadata["hidden_size"] == "64"
-    assert metadata["num_layers"] == "2"
+    assert metadata["hidden_size"] == 64
+    assert metadata["num_layers"] == 2
 
     # Test input/output details
     input_details = onnx_model.get_input_details()