kinfer 0.0.1__tar.gz → 0.0.2__tar.gz

{kinfer-0.0.1 → kinfer-0.0.2}/LICENSE

@@ -1,6 +1,6 @@
 MIT License
 
-Copyright (c) 2023 WT-MM
+Copyright (c) 2023 K-Scale Labs
 
 Permission is hereby granted, free of charge, to any person obtaining a copy
 of this software and associated documentation files (the "Software"), to deal

{kinfer-0.0.1/kinfer.egg-info → kinfer-0.0.2}/PKG-INFO

@@ -1,12 +1,15 @@
 Metadata-Version: 2.1
 Name: kinfer
-Version: 0.0.1
+Version: 0.0.2
 Summary: The kinfer project
 Home-page: https://github.com/kscalelabs/kinfer.git
-Author: WT-MM
-Requires-Python: >=3.11
+Author: K-Scale Labs
+Requires-Python: >=3.8
 Description-Content-Type: text/markdown
 License-File: LICENSE
+Requires-Dist: torch
+Requires-Dist: onnx
+Requires-Dist: onnxruntime
 Provides-Extra: dev
 Requires-Dist: black; extra == "dev"
 Requires-Dist: darglint; extra == "dev"

kinfer-0.0.2/kinfer/__init__.py

@@ -0,0 +1 @@
+__version__ = "0.0.2"

kinfer-0.0.2/kinfer/export/pytorch.py

@@ -0,0 +1,213 @@
+"""PyTorch model export utilities."""
+
+import inspect
+import sys
+from dataclasses import fields, is_dataclass
+from io import BytesIO
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
+
+import onnx
+import onnxruntime as ort  # type: ignore[import-untyped]
+import torch
+from torch import nn
+
+
+def get_model_info(model: nn.Module) -> Dict[str, Any]:
+    """Extract model information including input parameters and their types.
+
+    Args:
+        model: PyTorch model to analyze
+
+    Returns:
+        Dictionary containing model information
+    """
+    # Get model's forward method signature
+    signature = inspect.signature(model.forward)
+
+    # Extract parameter information
+    params_info = {}
+    for name, param in signature.parameters.items():
+        if name == "self":
+            continue
+        params_info[name] = {
+            "annotation": str(param.annotation),
+            "default": None if param.default is param.empty else str(param.default),
+        }
+
+    return {
+        "input_params": params_info,
+        "num_parameters": sum(p.numel() for p in model.parameters()),
+    }
+
+
+def add_metadata_to_onnx(
+    model_proto: onnx.ModelProto, metadata: Dict[str, Any], config: Optional[object] = None
+) -> onnx.ModelProto:
+    """Add metadata to ONNX model.
+
+    Args:
+        model_proto: ONNX model prototype
+        metadata: Dictionary of metadata to add
+        config: Optional configuration dataclass to add to metadata
+
+    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)
+
+    # 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)
+
+    return model_proto
+
+
+def infer_input_shapes(model: nn.Module) -> Union[torch.Size, List[torch.Size]]:
+    """Infer input shapes from model architecture.
+
+    Args:
+        model: PyTorch model to analyze
+
+    Returns:
+        Single input shape or list of input shapes
+    """
+    # Check if model is Sequential or has Sequential as first child
+    if isinstance(model, nn.Sequential):
+        first_layer = model[0]
+    else:
+        # Get the first immediate child
+        children = list(model.children())
+        first_layer = children[0] if children else None
+
+        # Unwrap if the first child is Sequential
+        if isinstance(first_layer, nn.Sequential):
+            first_layer = first_layer[0]
+    # Check if first layer is a type we can infer from
+    if not isinstance(first_layer, (nn.Linear, nn.Conv1d, nn.Conv2d, nn.Conv3d)):
+        raise ValueError("First layer must be Linear or Conv layer to infer input shape")
+
+    # Get input dimensions
+    if isinstance(first_layer, nn.Linear):
+        return torch.Size([1, first_layer.in_features])
+    elif isinstance(first_layer, nn.Conv1d):
+        raise ValueError("Cannot infer sequence length for Conv1d layer. Please provide input_tensors explicitly.")
+    elif isinstance(first_layer, nn.Conv2d):
+        raise ValueError("Cannot infer image dimensions for Conv2d layer. Please provide input_tensors explicitly.")
+    elif isinstance(first_layer, nn.Conv3d):
+        raise ValueError("Cannot infer volume dimensions for Conv3d layer. Please provide input_tensors explicitly.")
+
+    raise ValueError("Could not infer input shape from model architecture")
+
+
+def create_example_inputs(model: nn.Module) -> Union[torch.Tensor, Tuple[torch.Tensor, ...]]:
+    """Create example input tensors based on model's forward signature and architecture.
+
+    Args:
+        model: PyTorch model to analyze
+
+    Returns:
+        Single tensor or tuple of tensors matching the model's expected input
+    """
+    signature = inspect.signature(model.forward)
+    params = [p for p in signature.parameters.items() if p[0] != "self"]
+
+    # If single parameter (besides self), try to infer shape
+    if len(params) == 1:
+        shape = infer_input_shapes(model)
+        return torch.randn(*shape) if isinstance(shape, torch.Size) else torch.randn(*shape[0])
+
+    # For multiple parameters, try to infer from parameter annotations
+    input_tensors = []
+    for name, param in params:
+        # Try to get shape from annotation
+        if hasattr(param.annotation, "__origin__") and param.annotation.__origin__ is torch.Tensor:
+            # If annotation includes size information (e.g., Tensor[batch_size, channels, height, width])
+            if hasattr(param.annotation, "__args__"):
+                shape = param.annotation.__args__
+                input_tensors.append(torch.randn(*shape) if isinstance(shape, torch.Size) else torch.randn(*shape[0]))
+            else:
+                # Default to a vector if no size info
+                input_tensors.append(torch.randn(1, 32))
+        else:
+            # Default fallback
+            input_tensors.append(torch.randn(1, 32))
+
+    return tuple(input_tensors)
+
+
+def export_to_onnx(
+    model: nn.Module,
+    input_tensors: Optional[Union[torch.Tensor, Tuple[torch.Tensor, ...]]] = None,
+    config: Optional[object] = None,
+    save_path: Optional[str] = None,
+) -> ort.InferenceSession:
+    """Export PyTorch model to ONNX format with metadata.
+
+    Args:
+        model: PyTorch model to export
+        input_tensors: Optional example input tensors for model tracing. If None, will attempt to infer.
+        config: Optional configuration dataclass to add to metadata
+        save_path: Optional path to save the ONNX model
+
+    Returns:
+        ONNX inference session
+    """
+    # Get model information
+    model_info = get_model_info(model)
+
+    # Create example inputs if not provided
+    if input_tensors is None:
+        try:
+            input_tensors = create_example_inputs(model)
+            model_info["inferred_input_shapes"] = str(
+                input_tensors.shape if isinstance(input_tensors, torch.Tensor) else [t.shape for t in input_tensors]
+            )
+        except ValueError as e:
+            raise ValueError(
+                f"Could not automatically infer input shapes. Please provide input_tensors. Error: {str(e)}"
+            )
+
+    # Convert model to JIT if not already
+    if not isinstance(model, torch.jit.ScriptModule):
+        model = torch.jit.script(model)
+
+    # Export model to buffer
+    buffer = BytesIO()
+    if TYPE_CHECKING and sys.version_info >= (3, 11):
+        torch.onnx.export(
+            model,
+            (input_tensors,) if isinstance(input_tensors, torch.Tensor) else input_tensors,
+            buffer,  # type: ignore[arg-type]
+        )
+    else:
+        torch.onnx.export(
+            model,
+            (input_tensors,) if isinstance(input_tensors, torch.Tensor) else input_tensors,
+            buffer,
+        )
+    buffer.seek(0)
+
+    # Load as ONNX model
+    model_proto = onnx.load_model(buffer)
+
+    # Add metadata
+    model_proto = add_metadata_to_onnx(model_proto, model_info, config)
+
+    # Save if path provided
+    if save_path:
+        onnx.save_model(model_proto, save_path)
+
+    # Convert to inference session
+    buffer = BytesIO()
+    onnx.save_model(model_proto, buffer)
+    buffer.seek(0)
+
+    return ort.InferenceSession(buffer.read())

kinfer-0.0.2/kinfer/inference/python.py

@@ -0,0 +1,91 @@
+"""ONNX model inference utilities for Python."""
+
+from typing import Any, Dict, List, Union
+
+import numpy as np
+import onnx
+import onnxruntime as ort  # type: ignore[import-untyped]
+
+
+class ONNXModel:
+    """Wrapper for ONNX model inference."""
+
+    def __init__(
+        self,
+        model_path: str,
+    ) -> None:
+        """Initialize ONNX model.
+
+        Args:
+            model_path: Path to ONNX model file
+            config: Optional inference configuration
+        """
+        self.model_path = model_path
+
+        # Load model and create inference session
+        self.model = onnx.load(model_path)
+        self.session = ort.InferenceSession(
+            model_path,
+        )
+
+        # Extract metadata
+        self.metadata = {prop.key: prop.value for prop in self.model.metadata_props}
+
+        # Get input and output details
+        self.input_details = [{"name": x.name, "shape": x.shape, "type": x.type} for x in self.session.get_inputs()]
+        self.output_details = [{"name": x.name, "shape": x.shape, "type": x.type} for x in self.session.get_outputs()]
+
+    def __call__(
+        self, inputs: Union[np.ndarray, Dict[str, np.ndarray], List[np.ndarray]]
+    ) -> Union[np.ndarray, Dict[str, np.ndarray], List[np.ndarray]]:
+        """Run inference on input data.
+
+        Args:
+            inputs: Input data as numpy array, dictionary of arrays, or list of arrays
+
+        Returns:
+            Model outputs in the same format as inputs
+        """
+        # Convert single array to dict
+        if isinstance(inputs, np.ndarray):
+            input_dict = {self.input_details[0]["name"]: inputs}
+        # Convert list to dict
+        elif isinstance(inputs, list):
+            input_dict = {detail["name"]: arr for detail, arr in zip(self.input_details, inputs)}
+        else:
+            input_dict = inputs
+
+        # Run inference - pass None to output_names param to get all outputs
+        outputs = self.session.run(None, input_dict)
+
+        # Convert output format to match input
+        if isinstance(inputs, np.ndarray):
+            return outputs[0]
+        elif isinstance(inputs, list):
+            return outputs
+        else:
+            return {detail["name"]: arr for detail, arr in zip(self.output_details, outputs)}
+
+    def get_metadata(self) -> Dict[str, str]:
+        """Get model metadata.
+
+        Returns:
+            Dictionary of metadata key-value pairs
+        """
+        return self.metadata
+
+    def get_input_details(self) -> List[Dict[str, Any]]:
+        """Get input tensor details.
+
+        Returns:
+            List of dictionaries containing input tensor information
+        """
+        return self.input_details
+
+    def get_output_details(self) -> List[Dict[str, Any]]:
+        """Get output tensor details.
+
+        Returns:
+            List of dictionaries containing output tensor information
+        """
+        return self.output_details

kinfer-0.0.2/kinfer/requirements.txt

@@ -0,0 +1,5 @@
+# requirements.txt
+
+torch
+onnx
+onnxruntime

{kinfer-0.0.1 → kinfer-0.0.2/kinfer.egg-info}/PKG-INFO

@@ -1,12 +1,15 @@
 Metadata-Version: 2.1
 Name: kinfer
-Version: 0.0.1
+Version: 0.0.2
 Summary: The kinfer project
 Home-page: https://github.com/kscalelabs/kinfer.git
-Author: WT-MM
-Requires-Python: >=3.11
+Author: K-Scale Labs
+Requires-Python: >=3.8
 Description-Content-Type: text/markdown
 License-File: LICENSE
+Requires-Dist: torch
+Requires-Dist: onnx
+Requires-Dist: onnxruntime
 Provides-Extra: dev
 Requires-Dist: black; extra == "dev"
 Requires-Dist: darglint; extra == "dev"

{kinfer-0.0.1 → kinfer-0.0.2}/kinfer.egg-info/SOURCES.txt

@@ -12,4 +12,9 @@ kinfer.egg-info/SOURCES.txt
 kinfer.egg-info/dependency_links.txt
 kinfer.egg-info/requires.txt
 kinfer.egg-info/top_level.txt
-tests/test_dummy.py
+kinfer/export/__init__.py
+kinfer/export/pytorch.py
+kinfer/inference/__init__.py
+kinfer/inference/python.py
+tests/test_dummy.py
+tests/test_infer.py

{kinfer-0.0.1 → kinfer-0.0.2}/kinfer.egg-info/requires.txt

@@ -1,3 +1,6 @@
+torch
+onnx
+onnxruntime
 
 [dev]
 black

{kinfer-0.0.1 → kinfer-0.0.2}/pyproject.toml

@@ -46,7 +46,7 @@ target-version = "py310"
 
 [tool.ruff.lint]
 
-select = ["ANN", "D", "E", "F", "I", "N", "PGH", "PLC", "PLE", "PLR", "PLW", "W"]
+select = ["ANN", "D", "E", "F", "G", "I", "N", "PGH", "PLC", "PLE", "PLR", "PLW", "W"]
 
 ignore = [
     "ANN101", "ANN102",

{kinfer-0.0.1 → kinfer-0.0.2}/setup.py

@@ -3,6 +3,7 @@
 """Setup script for the project."""
 
 import re
+from typing import List
 
 from setuptools import setup
 
@@ -11,11 +12,11 @@ with open("README.md", "r", encoding="utf-8") as f:
 
 
 with open("kinfer/requirements.txt", "r", encoding="utf-8") as f:
-    requirements: list[str] = f.read().splitlines()
+    requirements: List[str] = f.read().splitlines()
 
 
 with open("kinfer/requirements-dev.txt", "r", encoding="utf-8") as f:
-    requirements_dev: list[str] = f.read().splitlines()
+    requirements_dev: List[str] = f.read().splitlines()
 
 
 with open("kinfer/__init__.py", "r", encoding="utf-8") as fh:
@@ -28,11 +29,11 @@ setup(
     name="kinfer",
     version=version,
     description="The kinfer project",
-    author="WT-MM",
+    author="K-Scale Labs",
     url="https://github.com/kscalelabs/kinfer.git",
     long_description=long_description,
     long_description_content_type="text/markdown",
-    python_requires=">=3.11",
+    python_requires=">=3.8",
     install_requires=requirements,
     tests_require=requirements_dev,
     extras_require={"dev": requirements_dev},

kinfer-0.0.2/tests/test_infer.py

@@ -0,0 +1,193 @@
+"""Tests for model inference functionality."""
+
+from dataclasses import dataclass
+from pathlib import Path
+
+import numpy as np
+import pytest
+import torch
+
+from kinfer.export.pytorch import export_to_onnx
+from kinfer.inference.python import ONNXModel
+
+
+@dataclass
+class ModelConfig:
+    hidden_size: int = 64
+    num_layers: int = 2
+
+
+class SimpleModel(torch.nn.Module):
+    """A simple neural network model for demonstration."""
+
+    def __init__(self, config: ModelConfig) -> None:
+        super().__init__()
+        layers = []
+        in_features = 10
+
+        for _ in range(config.num_layers):
+            layers.extend([torch.nn.Linear(in_features, config.hidden_size), torch.nn.ReLU()])
+            in_features = config.hidden_size
+
+        layers.append(torch.nn.Linear(config.hidden_size, 1))
+        self.net = torch.nn.Sequential(*layers)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.net(x)
+
+
+@pytest.fixture
+def model_path(tmp_path: Path) -> str:
+    """Create and export a test model."""
+    # Create and export model
+    config = ModelConfig()
+    model = SimpleModel(config)
+
+    save_path = str(tmp_path / "test_model.onnx")
+    export_to_onnx(model=model, input_tensors=torch.randn(1, 10), config=config, save_path=save_path)
+
+    return save_path
+
+
+def test_model_loading(model_path: str) -> None:
+    """Test basic model loading functionality."""
+    # Test with default config
+    model = ONNXModel(model_path)
+    assert model is not None
+
+    model = ONNXModel(model_path)
+    assert model is not None
+
+
+def test_model_metadata(model_path: str) -> None:
+    """Test model metadata extraction."""
+    model = ONNXModel(model_path)
+    metadata = model.get_metadata()
+
+    # 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"
+
+
+def test_model_inference(model_path: str) -> None:
+    """Test model inference with different input formats."""
+    model = ONNXModel(model_path)
+
+    # Test with numpy array
+    input_data = np.random.randn(1, 10).astype(np.float32)
+    output = model(input_data)
+    assert isinstance(output, np.ndarray)
+    assert output.shape == (1, 1)
+
+    # Test with dictionary input
+    input_name = model.get_input_details()[0]["name"]
+    output = model({input_name: input_data})
+    assert isinstance(output, dict)
+
+    # Test with list input
+    output = model([input_data])
+    assert isinstance(output, list)
+
+
+def test_model_details(model_path: str) -> None:
+    """Test input/output detail extraction."""
+    model = ONNXModel(model_path)
+
+    # Check input details
+    input_details = model.get_input_details()
+    assert len(input_details) == 1
+    assert input_details[0]["shape"] == [1, 10]
+
+    # Check output details
+    output_details = model.get_output_details()
+    assert len(output_details) == 1
+    assert output_details[0]["shape"] == [1, 1]
+
+
+def test_comprehensive_model_workflow(tmp_path: Path) -> None:
+    """Test complete model workflow including export, loading and inference."""
+    # Create and export model
+    config = ModelConfig(hidden_size=64, num_layers=2)
+    model = SimpleModel(config)
+    input_tensor = torch.randn(1, 10)
+
+    save_path = str(tmp_path / "test_model.onnx")
+    export_to_onnx(model=model, input_tensors=input_tensor, config=config, save_path=save_path)
+
+    # Load model for inference
+    onnx_model = ONNXModel(save_path)
+
+    # Test metadata
+    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"
+
+    # Test input/output details
+    input_details = onnx_model.get_input_details()
+    assert len(input_details) == 1
+    assert input_details[0]["shape"] == [1, 10]
+
+    output_details = onnx_model.get_output_details()
+    assert len(output_details) == 1
+    assert output_details[0]["shape"] == [1, 1]
+
+    # Test inference with different input methods
+    input_data = np.random.randn(1, 10).astype(np.float32)
+
+    # Method 1: Direct numpy array input
+    output1 = onnx_model(input_data)
+    assert isinstance(output1, np.ndarray)
+    assert output1.shape == (1, 1)
+
+    # Method 2: Dictionary input
+    input_name = onnx_model.get_input_details()[0]["name"]
+    output2 = onnx_model({input_name: input_data})
+    assert isinstance(output2, dict)
+    assert len(output2) == 1
+    assert list(output2.values())[0].shape == (1, 1)
+
+    # Method 3: List input
+    output3 = onnx_model([input_data])
+    assert isinstance(output3, list)
+    assert len(output3) == 1
+    assert output3[0].shape == (1, 1)
+
+
+def test_export_with_given_input(tmp_path: Path) -> None:
+    """Test model export with explicitly provided input tensor."""
+    config = ModelConfig()
+    model = SimpleModel(config)
+
+    # Create specific input tensor
+    input_tensor = torch.randn(1, 10)
+
+    save_path = str(tmp_path / "explicit_input_model.onnx")
+    session = export_to_onnx(model=model, input_tensors=input_tensor, config=config, save_path=save_path)
+
+    # Verify input shape matches what we provided
+    inputs = session.get_inputs()
+    assert len(inputs) == 1
+    assert inputs[0].shape == [1, 10]
+
+
+def test_export_with_inferred_input(tmp_path: Path) -> None:
+    """Test model export with automatically inferred input tensor."""
+    config = ModelConfig()
+    model = SimpleModel(config)
+
+    save_path = str(tmp_path / "inferred_input_model.onnx")
+    session = export_to_onnx(
+        model=model,
+        input_tensors=None,
+        config=config,
+        save_path=save_path,  # Let it infer the input
+    )
+
+    # Verify input shape was correctly inferred
+    inputs = session.get_inputs()
+    assert len(inputs) == 1
+    assert inputs[0].shape == [1, 10]  # Should match the in_features=10 from SimpleModel

kinfer-0.0.1/kinfer/__init__.py

@@ -1 +0,0 @@
-__version__ = "0.0.1"

kinfer-0.0.1/kinfer/requirements.txt

@@ -1 +0,0 @@
-# requirements.txt