haoline 0.3.0__py3-none-any.whl

haoline/tests/test_pytorch.py

@@ -0,0 +1,327 @@
+# Copyright (c) 2025 HaoLine Contributors
+# SPDX-License-Identifier: MIT
+
+"""
+Unit tests for PyTorch to ONNX conversion functionality.
+"""
+
+from __future__ import annotations
+
+import logging
+import sys
+from pathlib import Path
+from unittest.mock import patch
+
+import pytest
+
+# Add parent path for test imports
+sys.path.insert(0, str(Path(__file__).parent.parent.parent.parent))
+
+from ..cli import (
+    _convert_pytorch_to_onnx,
+    _extract_ultralytics_metadata,
+)
+from ..report import DatasetInfo, infer_num_classes_from_output
+
+# Check if torch is available
+try:
+    import torch
+    import torch.nn as nn
+
+    _TORCH_AVAILABLE = True
+except ImportError:
+    _TORCH_AVAILABLE = False
+    nn = None  # Placeholder for type hints
+
+
+# Only define the model class if torch is available
+if _TORCH_AVAILABLE:
+
+    class SimpleTestModel(nn.Module):
+        """Simple model for testing conversion."""
+
+        def __init__(self):
+            super().__init__()
+            self.conv = nn.Conv2d(3, 16, 3, padding=1)
+            self.relu = nn.ReLU()
+            self.pool = nn.AdaptiveAvgPool2d(1)
+            self.fc = nn.Linear(16, 10)
+
+        def forward(self, x):
+            x = self.relu(self.conv(x))
+            x = self.pool(x)
+            x = x.view(x.size(0), -1)
+            return self.fc(x)
+
+
+@pytest.mark.skipif(not _TORCH_AVAILABLE, reason="PyTorch not installed")
+class TestPyTorchConversion:
+    """Tests for PyTorch to ONNX conversion."""
+
+    def test_torchscript_model_conversion(self, tmp_path):
+        """TorchScript models should convert successfully."""
+        # Create and save a TorchScript model
+        model = SimpleTestModel()
+        model.eval()
+        dummy_input = torch.randn(1, 3, 32, 32)
+        traced = torch.jit.trace(model, dummy_input)
+
+        pt_path = tmp_path / "model.pt"
+        torch.jit.save(traced, str(pt_path))
+
+        logger = logging.getLogger("test")
+
+        # Convert
+        onnx_path, _temp_file = _convert_pytorch_to_onnx(
+            pt_path,
+            input_shape_str="1,3,32,32",
+            output_path=tmp_path / "output.onnx",
+            opset_version=17,
+            logger=logger,
+        )
+
+        assert onnx_path is not None
+        assert onnx_path.exists()
+        assert onnx_path.suffix == ".onnx"
+
+    def test_conversion_requires_input_shape(self, tmp_path):
+        """Conversion should fail without input shape."""
+        model = SimpleTestModel()
+        model.eval()
+        dummy_input = torch.randn(1, 3, 32, 32)
+        traced = torch.jit.trace(model, dummy_input)
+
+        pt_path = tmp_path / "model.pt"
+        torch.jit.save(traced, str(pt_path))
+
+        logger = logging.getLogger("test")
+
+        # Convert without input shape
+        onnx_path, _ = _convert_pytorch_to_onnx(
+            pt_path,
+            input_shape_str=None,  # No input shape
+            output_path=None,
+            opset_version=17,
+            logger=logger,
+        )
+
+        assert onnx_path is None
+
+    def test_conversion_invalid_input_shape(self, tmp_path):
+        """Conversion should fail with invalid input shape format."""
+        model = SimpleTestModel()
+        model.eval()
+        dummy_input = torch.randn(1, 3, 32, 32)
+        traced = torch.jit.trace(model, dummy_input)
+
+        pt_path = tmp_path / "model.pt"
+        torch.jit.save(traced, str(pt_path))
+
+        logger = logging.getLogger("test")
+
+        # Convert with invalid input shape
+        onnx_path, _ = _convert_pytorch_to_onnx(
+            pt_path,
+            input_shape_str="invalid,shape",
+            output_path=None,
+            opset_version=17,
+            logger=logger,
+        )
+
+        assert onnx_path is None
+
+    def test_conversion_nonexistent_file(self, tmp_path):
+        """Conversion should fail gracefully for nonexistent file."""
+        logger = logging.getLogger("test")
+
+        onnx_path, _ = _convert_pytorch_to_onnx(
+            tmp_path / "nonexistent.pt",
+            input_shape_str="1,3,32,32",
+            output_path=None,
+            opset_version=17,
+            logger=logger,
+        )
+
+        assert onnx_path is None
+
+    def test_conversion_temp_file_cleanup(self, tmp_path):
+        """Temp file should be created when no output path specified."""
+        model = SimpleTestModel()
+        model.eval()
+        dummy_input = torch.randn(1, 3, 32, 32)
+        traced = torch.jit.trace(model, dummy_input)
+
+        pt_path = tmp_path / "model.pt"
+        torch.jit.save(traced, str(pt_path))
+
+        logger = logging.getLogger("test")
+
+        # Convert without output path (should create temp file)
+        onnx_path, temp_file = _convert_pytorch_to_onnx(
+            pt_path,
+            input_shape_str="1,3,32,32",
+            output_path=None,  # No output path
+            opset_version=17,
+            logger=logger,
+        )
+
+        assert onnx_path is not None
+        assert temp_file is not None
+        assert onnx_path.exists()
+
+        # Clean up
+        onnx_path.unlink()
+
+    def test_state_dict_not_supported(self, tmp_path):
+        """State dict models should fail with helpful error."""
+        model = SimpleTestModel()
+
+        # Save as state_dict (not full model)
+        pt_path = tmp_path / "weights.pth"
+        torch.save(model.state_dict(), pt_path)
+
+        logger = logging.getLogger("test")
+
+        onnx_path, _ = _convert_pytorch_to_onnx(
+            pt_path,
+            input_shape_str="1,3,32,32",
+            output_path=None,
+            opset_version=17,
+            logger=logger,
+        )
+
+        assert onnx_path is None
+
+
+@pytest.mark.skipif(not _TORCH_AVAILABLE, reason="PyTorch not installed")
+class TestUltralyticsMetadataExtraction:
+    """Tests for Ultralytics metadata extraction."""
+
+    def test_extraction_without_ultralytics(self, tmp_path):
+        """Should return None gracefully when ultralytics not available."""
+        logger = logging.getLogger("test")
+
+        # Mock ultralytics not being available
+        with patch.dict("sys.modules", {"ultralytics": None}):
+            result = _extract_ultralytics_metadata(tmp_path / "fake.pt", logger)
+            # Should return None, not crash
+            assert result is None or isinstance(result, dict)
+
+
+class TestDatasetInfo:
+    """Tests for DatasetInfo dataclass."""
+
+    def test_dataset_info_creation(self):
+        """DatasetInfo should be created with expected fields."""
+        info = DatasetInfo(
+            task="detect",
+            num_classes=5,
+            class_names=["cat", "dog", "bird", "fish", "car"],
+            source="ultralytics",
+        )
+
+        assert info.task == "detect"
+        assert info.num_classes == 5
+        assert len(info.class_names) == 5
+        assert info.source == "ultralytics"
+
+    def test_dataset_info_defaults(self):
+        """DatasetInfo should have sensible defaults."""
+        info = DatasetInfo()
+
+        assert info.task is None
+        assert info.num_classes is None
+        assert info.class_names == []
+        assert info.source is None
+
+
+class TestInferNumClassesFromOutput:
+    """Tests for infer_num_classes_from_output function (Task 4B.2.2)."""
+
+    def test_classification_2d_output(self):
+        """Should detect classification from [batch, num_classes] shape."""
+        output_shapes = {"output": [1, 1000]}  # ImageNet-style
+        result = infer_num_classes_from_output(output_shapes)
+
+        assert result is not None
+        assert result.task == "classify"
+        assert result.num_classes == 1000
+        assert result.source == "output_shape"
+
+    def test_classification_3d_output(self):
+        """Should detect classification from [batch, 1, num_classes] shape."""
+        output_shapes = {"logits": [1, 1, 100]}  # CIFAR-100 style
+        result = infer_num_classes_from_output(output_shapes)
+
+        assert result is not None
+        assert result.task == "classify"
+        assert result.num_classes == 100
+        assert result.source == "output_shape"
+
+    def test_detection_yolo_output(self):
+        """Should detect detection from YOLO-style [batch, boxes, 4+nc] shape."""
+        # YOLOv8 output: [1, 8400, 84] for 80 COCO classes + 4 box coords
+        output_shapes = {"output0": [1, 8400, 84]}
+        result = infer_num_classes_from_output(output_shapes)
+
+        assert result is not None
+        assert result.task == "detect"
+        assert result.num_classes == 80  # 84 - 4 = 80
+        assert result.source == "output_shape"
+
+    def test_segmentation_output(self):
+        """Should detect segmentation from [batch, num_classes, h, w] shape."""
+        output_shapes = {"output": [1, 21, 512, 512]}  # Pascal VOC style
+        result = infer_num_classes_from_output(output_shapes)
+
+        assert result is not None
+        assert result.task == "segment"
+        assert result.num_classes == 21
+        assert result.source == "output_shape"
+
+    def test_empty_output_shapes(self):
+        """Should return None for empty output shapes."""
+        result = infer_num_classes_from_output({})
+        assert result is None
+
+    def test_single_output_dimension(self):
+        """Should return None for single-dimension outputs."""
+        output_shapes = {"output": [10]}
+        result = infer_num_classes_from_output(output_shapes)
+        assert result is None
+
+    def test_symbolic_dimensions(self):
+        """Should handle symbolic dimensions gracefully."""
+        output_shapes = {"output": ["batch", 1000]}
+        result = infer_num_classes_from_output(output_shapes)
+
+        assert result is not None
+        assert result.task == "classify"
+        assert result.num_classes == 1000
+
+    def test_priority_output_names(self):
+        """Should prioritize outputs with known names like 'logits'."""
+        output_shapes = {
+            "some_random_output": [1, 5],  # Would infer 5 classes
+            "logits": [1, 100],  # Should prefer this
+        }
+        result = infer_num_classes_from_output(output_shapes)
+
+        assert result is not None
+        assert result.num_classes == 100
+
+    def test_num_classes_too_small(self):
+        """Should not infer if num_classes is too small (< 2)."""
+        output_shapes = {"output": [1, 1]}  # Only 1 class - not valid
+        result = infer_num_classes_from_output(output_shapes)
+        assert result is None
+
+    def test_num_classes_too_large(self):
+        """Should not infer if num_classes is too large (> 10000)."""
+        output_shapes = {"output": [1, 50000]}  # Unlikely to be classes
+        result = infer_num_classes_from_output(output_shapes)
+        assert result is None
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])

haoline/tests/test_report.py

@@ -0,0 +1,383 @@
+# Copyright (c) 2025 HaoLine Contributors
+# SPDX-License-Identifier: MIT
+
+"""
+Unit tests for the report module (ModelInspector, InspectionReport).
+"""
+
+from __future__ import annotations
+
+import json
+import sys
+import tempfile
+from pathlib import Path
+
+import numpy as np
+import onnx
+import pytest
+from onnx import TensorProto, helper
+
+sys.path.insert(0, str(Path(__file__).parent.parent.parent.parent))
+from ..report import InspectionReport, ModelInspector
+
+
+def create_simple_model() -> onnx.ModelProto:
+    """Create a simple model for testing."""
+    X = helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 3, 8, 8])
+
+    W = helper.make_tensor(
+        "W",
+        TensorProto.FLOAT,
+        [16, 3, 3, 3],
+        np.random.randn(16, 3, 3, 3).astype(np.float32).flatten().tolist(),
+    )
+
+    Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1, 16, 6, 6])
+
+    conv = helper.make_node("Conv", ["X", "W"], ["conv_out"], kernel_shape=[3, 3])
+    relu = helper.make_node("Relu", ["conv_out"], ["Y"])
+
+    graph = helper.make_graph([conv, relu], "test", [X], [Y], [W])
+
+    model = helper.make_model(
+        graph,
+        opset_imports=[helper.make_opsetid("", 17)],
+        producer_name="test_producer",
+        producer_version="1.0",
+    )
+    return model
+
+
+class TestModelInspector:
+    """Tests for ModelInspector class."""
+
+    def test_inspect_returns_report(self):
+        """Inspect should return an InspectionReport."""
+        model = create_simple_model()
+
+        with tempfile.NamedTemporaryFile(suffix=".onnx", delete=False) as f:
+            onnx.save(model, f.name)
+            model_path = Path(f.name)
+
+        try:
+            inspector = ModelInspector()
+            report = inspector.inspect(model_path)
+
+            assert isinstance(report, InspectionReport)
+            assert report.metadata is not None
+            assert report.graph_summary is not None
+            assert report.param_counts is not None
+            assert report.flop_counts is not None
+            assert report.memory_estimates is not None
+        finally:
+            model_path.unlink()
+
+    def test_metadata_extraction(self):
+        """Test that metadata is extracted correctly."""
+        model = create_simple_model()
+
+        with tempfile.NamedTemporaryFile(suffix=".onnx", delete=False) as f:
+            onnx.save(model, f.name)
+            model_path = Path(f.name)
+
+        try:
+            inspector = ModelInspector()
+            report = inspector.inspect(model_path)
+
+            assert report.metadata.producer_name == "test_producer"
+            assert report.metadata.producer_version == "1.0"
+            assert "ai.onnx" in report.metadata.opsets
+        finally:
+            model_path.unlink()
+
+    def test_graph_summary(self):
+        """Test graph summary extraction."""
+        model = create_simple_model()
+
+        with tempfile.NamedTemporaryFile(suffix=".onnx", delete=False) as f:
+            onnx.save(model, f.name)
+            model_path = Path(f.name)
+
+        try:
+            inspector = ModelInspector()
+            report = inspector.inspect(model_path)
+
+            assert report.graph_summary.num_nodes == 2  # Conv + Relu
+            assert report.graph_summary.num_inputs == 1
+            assert report.graph_summary.num_outputs == 1
+            assert "Conv" in report.graph_summary.op_type_counts
+            assert "Relu" in report.graph_summary.op_type_counts
+        finally:
+            model_path.unlink()
+
+
+class TestInspectionReport:
+    """Tests for InspectionReport class."""
+
+    def test_to_json(self):
+        """Test JSON serialization."""
+        model = create_simple_model()
+
+        with tempfile.NamedTemporaryFile(suffix=".onnx", delete=False) as f:
+            onnx.save(model, f.name)
+            model_path = Path(f.name)
+
+        try:
+            inspector = ModelInspector()
+            report = inspector.inspect(model_path)
+
+            json_str = report.to_json()
+            assert json_str is not None
+
+            # Should be valid JSON
+            data = json.loads(json_str)
+            assert "metadata" in data
+            assert "graph_summary" in data
+            assert "param_counts" in data
+        finally:
+            model_path.unlink()
+
+    def test_to_dict(self):
+        """Test dictionary serialization."""
+        model = create_simple_model()
+
+        with tempfile.NamedTemporaryFile(suffix=".onnx", delete=False) as f:
+            onnx.save(model, f.name)
+            model_path = Path(f.name)
+
+        try:
+            inspector = ModelInspector()
+            report = inspector.inspect(model_path)
+
+            data = report.to_dict()
+            assert isinstance(data, dict)
+            assert "metadata" in data
+            assert data["metadata"]["producer_name"] == "test_producer"
+        finally:
+            model_path.unlink()
+
+    def test_to_markdown(self):
+        """Test Markdown generation."""
+        model = create_simple_model()
+
+        with tempfile.NamedTemporaryFile(suffix=".onnx", delete=False) as f:
+            onnx.save(model, f.name)
+            model_path = Path(f.name)
+
+        try:
+            inspector = ModelInspector()
+            report = inspector.inspect(model_path)
+
+            md = report.to_markdown()
+            assert isinstance(md, str)
+            assert "# Model Card:" in md
+            assert "## Metadata" in md
+            assert "## Graph Summary" in md
+        finally:
+            model_path.unlink()
+
+    def test_to_html_basic(self):
+        """Test HTML generation produces valid HTML structure."""
+        model = create_simple_model()
+
+        with tempfile.NamedTemporaryFile(suffix=".onnx", delete=False) as f:
+            onnx.save(model, f.name)
+            model_path = Path(f.name)
+
+        try:
+            inspector = ModelInspector()
+            report = inspector.inspect(model_path)
+
+            html = report.to_html()
+            assert isinstance(html, str)
+            assert "<!DOCTYPE html>" in html
+            assert "<html" in html
+            assert "</html>" in html
+            assert "HaoLine" in html
+        finally:
+            model_path.unlink()
+
+    def test_to_html_contains_sections(self):
+        """Test HTML contains expected sections."""
+        model = create_simple_model()
+
+        with tempfile.NamedTemporaryFile(suffix=".onnx", delete=False) as f:
+            onnx.save(model, f.name)
+            model_path = Path(f.name)
+
+        try:
+            inspector = ModelInspector()
+            report = inspector.inspect(model_path)
+
+            html = report.to_html()
+            # Should have key sections
+            assert "Model Details" in html
+            assert "Metadata" in html
+            assert "Graph Summary" in html
+        finally:
+            model_path.unlink()
+
+    def test_to_html_with_llm_summary(self):
+        """Test HTML includes LLM summary when provided."""
+        model = create_simple_model()
+
+        with tempfile.NamedTemporaryFile(suffix=".onnx", delete=False) as f:
+            onnx.save(model, f.name)
+            model_path = Path(f.name)
+
+        try:
+            inspector = ModelInspector()
+            report = inspector.inspect(model_path)
+
+            # Add mock LLM summary
+            report.llm_summary = {
+                "success": True,
+                "short_summary": "Test short summary.",
+                "detailed_summary": "Test detailed summary paragraph.",
+                "model": "test-model",
+            }
+
+            html = report.to_html()
+            assert "Executive Summary" in html
+            assert "Test short summary" in html
+            assert "Test detailed summary" in html
+        finally:
+            model_path.unlink()
+
+    def test_to_markdown_with_llm_summary(self):
+        """Test Markdown includes executive summary when LLM summary provided."""
+        model = create_simple_model()
+
+        with tempfile.NamedTemporaryFile(suffix=".onnx", delete=False) as f:
+            onnx.save(model, f.name)
+            model_path = Path(f.name)
+
+        try:
+            inspector = ModelInspector()
+            report = inspector.inspect(model_path)
+
+            # Add mock LLM summary
+            report.llm_summary = {
+                "success": True,
+                "short_summary": "This is a test model for unit testing.",
+                "detailed_summary": "This model contains Conv and ReLU operations for testing.",
+                "model": "test-llm-model",
+            }
+
+            md = report.to_markdown()
+            assert "## Executive Summary" in md
+            assert "Executive Summary" in md
+            assert "This is a test model for unit testing." in md
+            assert "This model contains Conv and ReLU operations for testing." in md
+            assert "*Generated by test-llm-model*" in md
+        finally:
+            model_path.unlink()
+
+    def test_to_markdown_no_executive_summary_without_llm(self):
+        """Test Markdown omits executive summary when no LLM summary."""
+        model = create_simple_model()
+
+        with tempfile.NamedTemporaryFile(suffix=".onnx", delete=False) as f:
+            onnx.save(model, f.name)
+            model_path = Path(f.name)
+
+        try:
+            inspector = ModelInspector()
+            report = inspector.inspect(model_path)
+
+            # No LLM summary set
+            md = report.to_markdown()
+            assert "## Executive Summary" not in md
+        finally:
+            model_path.unlink()
+
+    def test_to_html_embeds_images(self):
+        """Test HTML embeds images as base64 when provided."""
+        model = create_simple_model()
+
+        with tempfile.NamedTemporaryFile(suffix=".onnx", delete=False) as f:
+            onnx.save(model, f.name)
+            model_path = Path(f.name)
+
+        try:
+            inspector = ModelInspector()
+            report = inspector.inspect(model_path)
+
+            # Create a temp image file
+            with tempfile.TemporaryDirectory() as tmpdir:
+                # Create a minimal PNG (1x1 pixel)
+                import base64
+
+                # Minimal valid PNG
+                png_data = base64.b64decode(
+                    "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg=="
+                )
+                img_path = Path(tmpdir) / "test.png"
+                img_path.write_bytes(png_data)
+
+                html = report.to_html(image_paths={"test_chart": img_path})
+                assert "data:image/png;base64," in html
+        finally:
+            model_path.unlink()
+
+    def test_format_number(self):
+        """Test number formatting utility."""
+        assert InspectionReport._format_number(1_000) == "1.00K"
+        assert InspectionReport._format_number(1_000_000) == "1.00M"
+        assert InspectionReport._format_number(1_000_000_000) == "1.00B"
+        assert InspectionReport._format_number(500) == "500"
+
+    def test_format_bytes(self):
+        """Test byte formatting utility."""
+        assert InspectionReport._format_bytes(1024) == "1.02 KB"
+        assert InspectionReport._format_bytes(1024 * 1024) == "1.05 MB"
+        assert InspectionReport._format_bytes(1024 * 1024 * 1024) == "1.07 GB"
+        assert InspectionReport._format_bytes(500) == "500 bytes"
+
+
+class TestReportRiskSignals:
+    """Tests for risk signal inclusion in reports."""
+
+    def test_report_includes_risk_signals(self):
+        """Reports should include analyzed risk signals."""
+        model = create_simple_model()
+
+        with tempfile.NamedTemporaryFile(suffix=".onnx", delete=False) as f:
+            onnx.save(model, f.name)
+            model_path = Path(f.name)
+
+        try:
+            inspector = ModelInspector()
+            report = inspector.inspect(model_path)
+
+            # risk_signals is always a list (may be empty for small models)
+            assert isinstance(report.risk_signals, list)
+        finally:
+            model_path.unlink()
+
+    def test_report_includes_blocks(self):
+        """Reports should include detected blocks."""
+        model = create_simple_model()
+
+        with tempfile.NamedTemporaryFile(suffix=".onnx", delete=False) as f:
+            onnx.save(model, f.name)
+            model_path = Path(f.name)
+
+        try:
+            inspector = ModelInspector()
+            report = inspector.inspect(model_path)
+
+            assert isinstance(report.detected_blocks, list)
+            assert report.architecture_type in (
+                "cnn",
+                "mlp",
+                "transformer",
+                "hybrid",
+                "unknown",
+            )
+        finally:
+            model_path.unlink()
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])