haoline 0.3.0__py3-none-any.whl

haoline/pydantic_models.py

@@ -0,0 +1,241 @@
+# generated by datamodel-codegen:
+#   filename:  inspection_report_schema.json
+#   timestamp: 2025-12-04T22:53:59+00:00
+
+from __future__ import annotations
+
+from datetime import datetime
+from enum import Enum
+from typing import Annotated, Any
+
+from pydantic import BaseModel, Field
+
+
+class Metadata(BaseModel):
+    path: Annotated[str, Field(description="Path to the ONNX model file")]
+    ir_version: Annotated[int, Field(description="ONNX IR version", ge=1)]
+    producer_name: Annotated[str, Field(description="Name of the tool that produced the model")]
+    producer_version: Annotated[str | None, Field(description="Version of the producer tool")] = (
+        None
+    )
+    domain: Annotated[str | None, Field(description="Model domain")] = None
+    model_version: Annotated[int | None, Field(description="Model version number")] = None
+    doc_string: Annotated[str | None, Field(description="Model documentation string")] = None
+    opsets: Annotated[dict[str, int], Field(description="Opset versions by domain")]
+
+
+class GraphSummary(BaseModel):
+    num_nodes: Annotated[int | None, Field(description="Total number of nodes in graph", ge=0)] = (
+        None
+    )
+    num_inputs: Annotated[int | None, Field(description="Number of graph inputs", ge=0)] = None
+    num_outputs: Annotated[int | None, Field(description="Number of graph outputs", ge=0)] = None
+    num_initializers: Annotated[
+        int | None, Field(description="Number of initializers (weights)", ge=0)
+    ] = None
+    input_shapes: Annotated[
+        dict[str, list[int | str]] | None, Field(description="Input tensor shapes by name")
+    ] = None
+    output_shapes: Annotated[
+        dict[str, list[int | str]] | None,
+        Field(description="Output tensor shapes by name"),
+    ] = None
+    op_type_counts: Annotated[
+        dict[str, int] | None, Field(description="Count of each operator type")
+    ] = None
+
+
+class SharedWeights(BaseModel):
+    count: Annotated[
+        int | None, Field(description="Number of weights shared across 2+ nodes", ge=0)
+    ] = None
+    details: Annotated[
+        dict[str, list[str]] | None,
+        Field(description="Shared weight name to list of node names using it"),
+    ] = None
+
+
+class ParamCounts(BaseModel):
+    total: Annotated[int | None, Field(description="Total parameter count", ge=0)] = None
+    trainable: Annotated[int | None, Field(description="Trainable parameter count", ge=0)] = None
+    non_trainable: Annotated[
+        int | None, Field(description="Non-trainable parameter count", ge=0)
+    ] = None
+    by_op_type: Annotated[
+        dict[str, float] | None,
+        Field(description="Parameters by operator type (fractional for shared weights)"),
+    ] = None
+    shared_weights: Annotated[
+        SharedWeights | None, Field(description="Information about shared weights")
+    ] = None
+    precision_breakdown: Annotated[
+        dict[str, int] | None,
+        Field(description="Parameter count by data type (fp32, fp16, int8, etc.)"),
+    ] = None
+    is_quantized: Annotated[
+        bool | None, Field(description="Whether model uses quantized weights or ops")
+    ] = None
+    quantized_ops: Annotated[
+        list[str] | None, Field(description="List of quantized operation types detected")
+    ] = None
+
+
+class Hotspot(BaseModel):
+    name: str | None = None
+    op_type: str | None = None
+    flops: Annotated[int | None, Field(ge=0)] = None
+
+
+class FlopCounts(BaseModel):
+    total: Annotated[int | None, Field(description="Total estimated FLOPs", ge=0)] = None
+    by_node_type: Annotated[dict[str, int] | None, Field(description="FLOPs by operator type")] = (
+        None
+    )
+    hotspots: Annotated[list[Hotspot] | None, Field(description="Top compute-intensive nodes")] = (
+        None
+    )
+
+
+class KvCacheConfig(BaseModel):
+    num_layers: Annotated[int | None, Field(ge=0)] = None
+    hidden_dim: Annotated[int | None, Field(ge=0)] = None
+    seq_len: Annotated[int | None, Field(ge=0)] = None
+    bytes_per_element: Annotated[int | None, Field(ge=1)] = None
+
+
+class Breakdown(BaseModel):
+    weights_by_op_type: dict[str, int] | None = None
+    activations_by_op_type: dict[str, int] | None = None
+
+
+class MemoryEstimates(BaseModel):
+    model_size_bytes: Annotated[int | None, Field(description="Model size in bytes", ge=0)] = None
+    peak_activation_bytes: Annotated[
+        int | None, Field(description="Peak activation memory in bytes", ge=0)
+    ] = None
+    kv_cache_bytes_per_token: Annotated[
+        int | None, Field(description="KV cache memory per token (transformers)", ge=0)
+    ] = None
+    kv_cache_bytes_full_context: Annotated[
+        int | None, Field(description="KV cache for full context length", ge=0)
+    ] = None
+    kv_cache_config: Annotated[
+        KvCacheConfig | None, Field(description="KV cache configuration")
+    ] = None
+    breakdown: Annotated[Breakdown | None, Field(description="Memory breakdown by component")] = (
+        None
+    )
+
+
+class DetectedBlock(BaseModel):
+    block_type: Annotated[str, Field(description="Type of block (e.g., ResidualAdd, Attention)")]
+    name: Annotated[str, Field(description="Block identifier")]
+    nodes: Annotated[list[str] | None, Field(description="Node names in this block")] = None
+    start_node: str | None = None
+    end_node: str | None = None
+    attributes: Annotated[dict[str, Any] | None, Field(description="Block-specific attributes")] = (
+        None
+    )
+
+
+class ArchitectureType(Enum):
+    transformer = "transformer"
+    cnn = "cnn"
+    mlp = "mlp"
+    hybrid = "hybrid"
+    unknown = "unknown"
+
+
+class Severity(Enum):
+    info = "info"
+    warning = "warning"
+    high = "high"
+
+
+class RiskSignal(BaseModel):
+    id: Annotated[str, Field(description="Risk signal identifier")]
+    severity: Annotated[Severity, Field(description="Severity level")]
+    description: Annotated[str, Field(description="Human-readable description")]
+    nodes: Annotated[list[str] | None, Field(description="Affected node names")] = None
+    recommendation: Annotated[str | None, Field(description="Recommended action")] = None
+
+
+class HardwareProfile(BaseModel):
+    name: str | None = None
+    vram_bytes: Annotated[int | None, Field(ge=0)] = None
+    peak_fp32_tflops: Annotated[float | None, Field(ge=0.0)] = None
+    peak_fp16_tflops: Annotated[float | None, Field(ge=0.0)] = None
+    memory_bandwidth_gbps: Annotated[float | None, Field(ge=0.0)] = None
+    tdp_watts: int | None = None
+
+
+class HardwareEstimates(BaseModel):
+    device: str | None = None
+    precision: str | None = None
+    batch_size: Annotated[int | None, Field(ge=1)] = None
+    vram_required_bytes: Annotated[int | None, Field(ge=0)] = None
+    fits_in_vram: bool | None = None
+    theoretical_latency_ms: Annotated[float | None, Field(ge=0.0)] = None
+    bottleneck: str | None = None
+    compute_utilization_estimate: Annotated[float | None, Field(ge=0.0)] = None
+    gpu_saturation: Annotated[float | None, Field(ge=0.0)] = None
+
+
+class LlmSummary(BaseModel):
+    success: bool | None = None
+    short_summary: str | None = None
+    detailed_summary: str | None = None
+    model: str | None = None
+    error: str | None = None
+
+
+class DatasetInfo(BaseModel):
+    task: str | None = None
+    num_classes: Annotated[int | None, Field(ge=0)] = None
+    class_names: list[str] | None = None
+    source: str | None = None
+
+
+class HaolineInspectionReport(BaseModel):
+    metadata: Annotated[Metadata, Field(description="Model metadata extracted from ONNX proto")]
+    generated_at: Annotated[
+        datetime, Field(description="ISO 8601 timestamp when report was generated")
+    ]
+    autodoc_version: Annotated[
+        str,
+        Field(
+            description="Version of HaoLine that generated the report",
+            pattern=r"^[0-9]+\.[0-9]+\.[0-9]+",
+        ),
+    ]
+    graph_summary: Annotated[
+        GraphSummary | None, Field(description="Summary statistics about the ONNX graph")
+    ] = None
+    param_counts: Annotated[ParamCounts | None, Field(description="Parameter count statistics")] = (
+        None
+    )
+    flop_counts: Annotated[FlopCounts | None, Field(description="FLOP estimation statistics")] = (
+        None
+    )
+    memory_estimates: Annotated[
+        MemoryEstimates | None, Field(description="Memory usage estimates")
+    ] = None
+    detected_blocks: Annotated[
+        list[DetectedBlock] | None, Field(description="Detected architectural blocks")
+    ] = None
+    architecture_type: Annotated[
+        ArchitectureType | None, Field(description="Detected architecture type")
+    ] = None
+    risk_signals: Annotated[list[RiskSignal] | None, Field(description="Detected risk signals")] = (
+        None
+    )
+    hardware_profile: Annotated[
+        HardwareProfile | None, Field(description="Target hardware profile")
+    ] = None
+    hardware_estimates: Annotated[
+        HardwareEstimates | None, Field(description="Hardware-specific estimates")
+    ] = None
+    llm_summary: Annotated[LlmSummary | None, Field(description="LLM-generated summary")] = None
+    dataset_info: Annotated[
+        DatasetInfo | None, Field(description="Dataset and class information")
+    ] = None