haoline 0.3.0__py3-none-any.whl

haoline/formats/openvino.py

@@ -0,0 +1,255 @@
+# Copyright (c) 2025 HaoLine Contributors
+# SPDX-License-Identifier: MIT
+
+"""
+OpenVINO format reader.
+
+OpenVINO is Intel's toolkit for optimizing and deploying AI inference.
+This reader extracts model metadata from OpenVINO IR format (.xml/.bin).
+
+Requires: openvino (pip install openvino)
+
+Reference: https://docs.openvino.ai/
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any
+
+
+@dataclass
+class OpenVINOLayerInfo:
+    """Information about an OpenVINO layer."""
+
+    name: str
+    type: str
+    input_shapes: list[tuple[int, ...]] = field(default_factory=list)
+    output_shapes: list[tuple[int, ...]] = field(default_factory=list)
+    precision: str = "FP32"
+
+
+@dataclass
+class OpenVINOInfo:
+    """Parsed OpenVINO model information."""
+
+    path: Path
+    name: str
+    framework: str
+    layers: list[OpenVINOLayerInfo] = field(default_factory=list)
+    inputs: list[dict[str, Any]] = field(default_factory=list)
+    outputs: list[dict[str, Any]] = field(default_factory=list)
+    metadata: dict[str, Any] = field(default_factory=dict)
+
+    @property
+    def layer_count(self) -> int:
+        """Number of layers."""
+        return len(self.layers)
+
+    @property
+    def layer_type_counts(self) -> dict[str, int]:
+        """Count of layers by type."""
+        counts: dict[str, int] = {}
+        for layer in self.layers:
+            counts[layer.type] = counts.get(layer.type, 0) + 1
+        return counts
+
+    @property
+    def precision_breakdown(self) -> dict[str, int]:
+        """Count of layers by precision."""
+        breakdown: dict[str, int] = {}
+        for layer in self.layers:
+            breakdown[layer.precision] = breakdown.get(layer.precision, 0) + 1
+        return breakdown
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert to dictionary for JSON serialization."""
+        return {
+            "path": str(self.path),
+            "name": self.name,
+            "framework": self.framework,
+            "layer_count": self.layer_count,
+            "layer_type_counts": self.layer_type_counts,
+            "precision_breakdown": self.precision_breakdown,
+            "inputs": self.inputs,
+            "outputs": self.outputs,
+            "metadata": self.metadata,
+        }
+
+
+class OpenVINOReader:
+    """Reader for OpenVINO IR format (.xml/.bin)."""
+
+    def __init__(self, path: str | Path):
+        """
+        Initialize reader with file path.
+
+        Args:
+            path: Path to the OpenVINO model (.xml or directory).
+
+        Raises:
+            ImportError: If openvino is not installed.
+        """
+        self.path = Path(path)
+
+        # Handle both .xml path and directory containing model
+        if self.path.is_dir():
+            xml_files = list(self.path.glob("*.xml"))
+            if xml_files:
+                self.path = xml_files[0]
+            else:
+                raise FileNotFoundError(f"No .xml file found in: {self.path}")
+        elif self.path.suffix.lower() != ".xml":
+            # Try adding .xml
+            xml_path = self.path.with_suffix(".xml")
+            if xml_path.exists():
+                self.path = xml_path
+
+        if not self.path.exists():
+            raise FileNotFoundError(f"OpenVINO model not found: {self.path}")
+
+        try:
+            import openvino  # noqa: F401
+        except ImportError as e:
+            raise ImportError("openvino required. Install with: pip install openvino") from e
+
+    def read(self) -> OpenVINOInfo:
+        """
+        Read and parse the OpenVINO model.
+
+        Returns:
+            OpenVINOInfo with model metadata.
+        """
+        from openvino.runtime import Core
+
+        core = Core()
+        model = core.read_model(str(self.path))
+
+        # Extract basic info
+        name = model.get_friendly_name()
+
+        # Get RT info for framework
+        rt_info = model.get_rt_info()
+        framework = "unknown"
+        if "framework" in rt_info:
+            framework = str(rt_info["framework"].value)
+
+        # Extract inputs
+        inputs = []
+        for inp in model.inputs:
+            input_info = {
+                "name": inp.get_any_name(),
+                "shape": list(inp.get_partial_shape()),
+                "element_type": str(inp.get_element_type()),
+            }
+            inputs.append(input_info)
+
+        # Extract outputs
+        outputs = []
+        for out in model.outputs:
+            output_info = {
+                "name": out.get_any_name(),
+                "shape": list(out.get_partial_shape()),
+                "element_type": str(out.get_element_type()),
+            }
+            outputs.append(output_info)
+
+        # Extract layers (operations)
+        layers = []
+        for op in model.get_ordered_ops():
+            # Get input shapes
+            input_shapes = []
+            for inp in op.inputs():
+                shape = inp.get_partial_shape()
+                if shape.is_static:
+                    input_shapes.append(tuple(shape.to_shape()))
+                else:
+                    input_shapes.append(tuple(d.get_length() if d.is_static else -1 for d in shape))
+
+            # Get output shapes
+            output_shapes = []
+            for out in op.outputs():
+                shape = out.get_partial_shape()
+                if shape.is_static:
+                    output_shapes.append(tuple(shape.to_shape()))
+                else:
+                    output_shapes.append(
+                        tuple(d.get_length() if d.is_static else -1 for d in shape)
+                    )
+
+            # Get precision from output
+            precision = "FP32"
+            if op.outputs():
+                element_type = str(op.outputs()[0].get_element_type())
+                precision = element_type.upper().replace("F", "FP").replace("I", "INT")
+
+            layers.append(
+                OpenVINOLayerInfo(
+                    name=op.get_friendly_name(),
+                    type=op.get_type_name(),
+                    input_shapes=input_shapes,
+                    output_shapes=output_shapes,
+                    precision=precision,
+                )
+            )
+
+        # Additional metadata
+        metadata = {}
+        for key in rt_info:
+            try:
+                metadata[key] = str(rt_info[key].value)
+            except Exception:
+                pass
+
+        return OpenVINOInfo(
+            path=self.path,
+            name=name,
+            framework=framework,
+            layers=layers,
+            inputs=inputs,
+            outputs=outputs,
+            metadata=metadata,
+        )
+
+
+def is_openvino_file(path: str | Path) -> bool:
+    """
+    Check if a file is an OpenVINO model.
+
+    Args:
+        path: Path to check.
+
+    Returns:
+        True if the file is an OpenVINO IR model.
+    """
+    path = Path(path)
+    if not path.exists():
+        return False
+
+    # Check for .xml file
+    if path.suffix.lower() == ".xml":
+        # Quick check for OpenVINO XML structure
+        try:
+            with open(path, encoding="utf-8") as f:
+                header = f.read(500)
+                return "<net" in header and ("name=" in header or "version=" in header)
+        except Exception:
+            return False
+
+    # Check if .bin file has corresponding .xml
+    if path.suffix.lower() == ".bin":
+        xml_path = path.with_suffix(".xml")
+        return xml_path.exists() and is_openvino_file(xml_path)
+
+    return False
+
+
+def is_available() -> bool:
+    """Check if openvino is available."""
+    try:
+        import openvino  # noqa: F401
+
+        return True
+    except ImportError:
+        return False

haoline/formats/safetensors.py

@@ -0,0 +1,273 @@
+# Copyright (c) 2025 HaoLine Contributors
+# SPDX-License-Identifier: MIT
+
+"""
+SafeTensors format reader.
+
+SafeTensors is a simple, safe format for storing tensors, widely used
+in the HuggingFace ecosystem. Unlike pickle-based formats, SafeTensors
+cannot execute arbitrary code.
+
+This reader extracts:
+- Tensor names, shapes, and dtypes
+- Parameter counts and memory estimates
+- Metadata (if present)
+
+Reference: https://github.com/huggingface/safetensors
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any
+
+# Bytes per dtype
+DTYPE_SIZES: dict[str, int] = {
+    "F64": 8,
+    "F32": 4,
+    "F16": 2,
+    "BF16": 2,
+    "I64": 8,
+    "I32": 4,
+    "I16": 2,
+    "I8": 1,
+    "U8": 1,
+    "BOOL": 1,
+}
+
+
+@dataclass
+class SafeTensorInfo:
+    """Information about a single tensor."""
+
+    name: str
+    dtype: str
+    shape: tuple[int, ...]
+
+    @property
+    def n_elements(self) -> int:
+        """Total number of elements."""
+        result = 1
+        for d in self.shape:
+            result *= d
+        return result
+
+    @property
+    def size_bytes(self) -> int:
+        """Size in bytes."""
+        return self.n_elements * DTYPE_SIZES.get(self.dtype, 4)
+
+
+@dataclass
+class SafeTensorsInfo:
+    """Parsed SafeTensors file information."""
+
+    path: Path
+    tensors: list[SafeTensorInfo] = field(default_factory=list)
+    metadata: dict[str, str] = field(default_factory=dict)
+
+    @property
+    def total_params(self) -> int:
+        """Total parameter count."""
+        return sum(t.n_elements for t in self.tensors)
+
+    @property
+    def total_size_bytes(self) -> int:
+        """Total size in bytes."""
+        return sum(t.size_bytes for t in self.tensors)
+
+    @property
+    def dtype_breakdown(self) -> dict[str, int]:
+        """Count of tensors by dtype."""
+        breakdown: dict[str, int] = {}
+        for t in self.tensors:
+            breakdown[t.dtype] = breakdown.get(t.dtype, 0) + 1
+        return breakdown
+
+    @property
+    def size_breakdown(self) -> dict[str, int]:
+        """Size in bytes by dtype."""
+        breakdown: dict[str, int] = {}
+        for t in self.tensors:
+            breakdown[t.dtype] = breakdown.get(t.dtype, 0) + t.size_bytes
+        return breakdown
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert to dictionary for JSON serialization."""
+        return {
+            "path": str(self.path),
+            "tensor_count": len(self.tensors),
+            "total_params": self.total_params,
+            "total_size_bytes": self.total_size_bytes,
+            "dtype_breakdown": self.dtype_breakdown,
+            "size_breakdown": self.size_breakdown,
+            "metadata": self.metadata,
+            "tensors": [{"name": t.name, "dtype": t.dtype, "shape": t.shape} for t in self.tensors],
+        }
+
+
+class SafeTensorsReader:
+    """Reader for SafeTensors format files."""
+
+    def __init__(self, path: str | Path):
+        """
+        Initialize reader with file path.
+
+        Args:
+            path: Path to the SafeTensors file.
+
+        Raises:
+            ImportError: If safetensors library is not installed.
+        """
+        self.path = Path(path)
+        if not self.path.exists():
+            raise FileNotFoundError(f"SafeTensors file not found: {self.path}")
+
+        try:
+            import safetensors  # noqa: F401
+        except ImportError as e:
+            raise ImportError(
+                "safetensors library required. Install with: pip install safetensors"
+            ) from e
+
+    def read(self) -> SafeTensorsInfo:
+        """
+        Read and parse the SafeTensors file.
+
+        Returns:
+            SafeTensorsInfo with tensor information.
+        """
+        from safetensors import safe_open
+
+        tensors = []
+        metadata = {}
+
+        with safe_open(self.path, framework="np") as f:
+            # Get metadata if present
+            meta = f.metadata()
+            if meta:
+                metadata = dict(meta)
+
+            # Get tensor info
+            for name in f.keys():
+                tensor = f.get_tensor(name)
+                info = SafeTensorInfo(
+                    name=name,
+                    dtype=self._numpy_dtype_to_safetensors(tensor.dtype),
+                    shape=tuple(tensor.shape),
+                )
+                tensors.append(info)
+
+        return SafeTensorsInfo(
+            path=self.path,
+            tensors=tensors,
+            metadata=metadata,
+        )
+
+    def read_header_only(self) -> SafeTensorsInfo:
+        """
+        Read only the header without loading tensor data.
+
+        This is faster for large files when you only need metadata.
+
+        Returns:
+            SafeTensorsInfo with tensor information.
+        """
+        import json
+
+        tensors = []
+        metadata = {}
+
+        # SafeTensors header is at the start of the file
+        # Format: 8 bytes (header size as u64) + header JSON
+        with open(self.path, "rb") as f:
+            # Read header size
+            header_size = int.from_bytes(f.read(8), "little")
+
+            # Read header JSON
+            header_bytes = f.read(header_size)
+            header = json.loads(header_bytes)
+
+            # Extract metadata
+            if "__metadata__" in header:
+                metadata = header.pop("__metadata__")
+
+            # Extract tensor info
+            for name, info in header.items():
+                dtype = info["dtype"]
+                shape = tuple(info["shape"])
+                tensors.append(SafeTensorInfo(name=name, dtype=dtype, shape=shape))
+
+        return SafeTensorsInfo(
+            path=self.path,
+            tensors=tensors,
+            metadata=metadata,
+        )
+
+    def _numpy_dtype_to_safetensors(self, dtype) -> str:
+        """Convert numpy dtype to safetensors dtype string."""
+        import numpy as np
+
+        dtype_map = {
+            np.float64: "F64",
+            np.float32: "F32",
+            np.float16: "F16",
+            np.int64: "I64",
+            np.int32: "I32",
+            np.int16: "I16",
+            np.int8: "I8",
+            np.uint8: "U8",
+            np.bool_: "BOOL",
+        }
+
+        # Handle bfloat16 specially (might be stored as uint16)
+        dtype_name = str(dtype)
+        if "bfloat16" in dtype_name:
+            return "BF16"
+
+        return dtype_map.get(dtype.type, "F32")
+
+
+def is_safetensors_file(path: str | Path) -> bool:
+    """
+    Check if a file is a valid SafeTensors file.
+
+    Args:
+        path: Path to check.
+
+    Returns:
+        True if the file appears to be a valid SafeTensors file.
+    """
+    path = Path(path)
+    if not path.exists() or not path.is_file():
+        return False
+
+    # Check extension
+    if path.suffix.lower() != ".safetensors":
+        return False
+
+    # Try to read header
+    try:
+        with open(path, "rb") as f:
+            header_size = int.from_bytes(f.read(8), "little")
+            # Sanity check: header shouldn't be larger than 100MB
+            if header_size > 100 * 1024 * 1024:
+                return False
+            # Try to parse as JSON
+
+            header_bytes = f.read(min(header_size, 1024))  # Just read start
+            # Should start with '{'
+            return header_bytes.startswith(b"{")
+    except Exception:
+        return False
+
+
+def is_available() -> bool:
+    """Check if safetensors library is available."""
+    try:
+        import safetensors  # noqa: F401
+
+        return True
+    except ImportError:
+        return False