raxe 0.4.6__py3-none-any.whl

raxe/domain/ml/model_metadata.py

@@ -0,0 +1,263 @@
+"""Model metadata schema for L2 model registry.
+
+Defines the structure for model metadata files that describe
+available L2 models (ONNX, PyTorch, different versions, etc.).
+"""
+from dataclasses import dataclass, field
+from enum import Enum
+from pathlib import Path
+
+
+class ModelStatus(Enum):
+    """Model status for filtering."""
+    ACTIVE = "active"           # Production-ready, recommended
+    EXPERIMENTAL = "experimental"  # Testing, not for production
+    DEPRECATED = "deprecated"   # Old, being phased out
+
+
+class ModelRuntime(Enum):
+    """Model runtime type."""
+    PYTORCH = "pytorch"
+    ONNX = "onnx"
+    ONNX_INT8 = "onnx_int8"
+    ONNX_FP16 = "onnx_fp16"
+    TENSORFLOW = "tensorflow"
+    CUSTOM = "custom"
+
+
+@dataclass
+class FileInfo:
+    """Model file information."""
+    filename: str
+    size_mb: float = 0.0
+    checksum: str | None = None
+    onnx_embeddings: str | None = None  # Optional ONNX embeddings model path
+
+
+@dataclass
+class PerformanceMetrics:
+    """Performance characteristics."""
+    target_latency_ms: float
+    p50_latency_ms: float | None = None
+    p95_latency_ms: float | None = None
+    p99_latency_ms: float | None = None
+    throughput_per_sec: int | None = None
+    memory_mb: int | None = None
+
+
+@dataclass
+class AccuracyMetrics:
+    """Accuracy metrics."""
+    binary_f1: float | None = None
+    family_f1: float | None = None
+    subfamily_f1: float | None = None
+    false_positive_rate: float | None = None
+    false_negative_rate: float | None = None
+
+
+@dataclass
+class Requirements:
+    """Runtime requirements."""
+    runtime: ModelRuntime
+    min_runtime_version: str | None = None
+    requires_gpu: bool = False
+    requires_quantization_support: bool = False
+    additional_dependencies: list[str] = field(default_factory=list)
+
+
+@dataclass
+class ModelMetadata:
+    """Complete model metadata.
+
+    This describes a single L2 model variant with all its characteristics,
+    performance metrics, and requirements.
+
+    Example:
+        metadata = ModelMetadata(
+            model_id="v1.0_onnx_int8",
+            name="RAXE L2 v1.0 ONNX INT8",
+            version="0.0.1",
+            variant="onnx_int8",
+            description="Optimized ONNX INT8 model",
+            file_info=FileInfo(filename="raxe_model_l2_v1.0_onnx_int8.raxe"),
+            performance=PerformanceMetrics(target_latency_ms=10),
+            requirements=Requirements(runtime=ModelRuntime.ONNX_INT8),
+            status=ModelStatus.ACTIVE,
+            tokenizer_name="sentence-transformers/all-mpnet-base-v2",
+            tokenizer_config={"max_length": 512, "type": "AutoTokenizer"},
+            embedding_model_name="all-mpnet-base-v2"
+        )
+    """
+    model_id: str
+    name: str
+    version: str
+    variant: str
+    description: str
+
+    file_info: FileInfo
+    performance: PerformanceMetrics
+    requirements: Requirements
+
+    status: ModelStatus = ModelStatus.ACTIVE
+    accuracy: AccuracyMetrics | None = None
+    tags: list[str] = field(default_factory=list)
+    recommended_for: list[str] = field(default_factory=list)
+    not_recommended_for: list[str] = field(default_factory=list)
+
+    # Tokenizer fields (for manifest-based models)
+    tokenizer_name: str | None = None
+    tokenizer_config: dict | None = None
+    embedding_model_name: str | None = None
+
+    # Computed fields
+    file_path: Path | None = None
+
+    def __post_init__(self):
+        """Validate metadata."""
+        if not self.model_id:
+            raise ValueError("model_id is required")
+        if not self.file_info.filename:
+            raise ValueError("filename is required")
+
+    @property
+    def is_active(self) -> bool:
+        """True if model is active (production-ready)."""
+        return self.status == ModelStatus.ACTIVE
+
+    @property
+    def is_experimental(self) -> bool:
+        """True if model is experimental."""
+        return self.status == ModelStatus.EXPERIMENTAL
+
+    @property
+    def runtime_type(self) -> str:
+        """Runtime type as string."""
+        return self.requirements.runtime.value
+
+    def score_for_criteria(self, criteria: str) -> float:
+        """Calculate score for selection criteria.
+
+        Args:
+            criteria: "latency", "accuracy", "balanced", "memory"
+
+        Returns:
+            Score from 0-100 (higher is better)
+        """
+        if criteria == "latency":
+            # Lower latency is better
+            if self.performance.p95_latency_ms:
+                # Score: 100 at 5ms, 50 at 50ms, 0 at 100ms+
+                latency = self.performance.p95_latency_ms
+                return max(0, 100 - (latency - 5))
+            return 50  # Default if unknown
+
+        elif criteria == "accuracy":
+            # Higher accuracy is better
+            if self.accuracy and self.accuracy.binary_f1:
+                # Score: F1 * 100
+                return self.accuracy.binary_f1 * 100
+            return 50  # Default if unknown
+
+        elif criteria == "memory":
+            # Lower memory is better
+            if self.performance.memory_mb:
+                # Score: 100 at 50MB, 50 at 150MB, 0 at 250MB+
+                memory = self.performance.memory_mb
+                return max(0, 100 - (memory - 50) / 2)
+            return 50
+
+        elif criteria == "balanced":
+            # Weighted average of latency and accuracy
+            latency_score = self.score_for_criteria("latency")
+            accuracy_score = self.score_for_criteria("accuracy")
+            return (latency_score * 0.6) + (accuracy_score * 0.4)
+
+        else:
+            return 50
+
+    def to_dict(self) -> dict:
+        """Convert to dictionary for serialization."""
+        result = {
+            "model_id": self.model_id,
+            "name": self.name,
+            "version": self.version,
+            "variant": self.variant,
+            "description": self.description,
+            "file_info": {
+                "filename": self.file_info.filename,
+                "size_mb": self.file_info.size_mb,
+                "checksum": self.file_info.checksum,
+            },
+            "performance": {
+                "target_latency_ms": self.performance.target_latency_ms,
+                "p50_latency_ms": self.performance.p50_latency_ms,
+                "p95_latency_ms": self.performance.p95_latency_ms,
+                "p99_latency_ms": self.performance.p99_latency_ms,
+                "throughput_per_sec": self.performance.throughput_per_sec,
+                "memory_mb": self.performance.memory_mb,
+            },
+            "accuracy": {
+                "binary_f1": self.accuracy.binary_f1 if self.accuracy else None,
+                "family_f1": self.accuracy.family_f1 if self.accuracy else None,
+                "subfamily_f1": self.accuracy.subfamily_f1 if self.accuracy else None,
+            } if self.accuracy else None,
+            "requirements": {
+                "runtime": self.runtime_type,
+                "requires_gpu": self.requirements.requires_gpu,
+            },
+            "status": self.status.value,
+            "tags": self.tags,
+        }
+
+        # Add tokenizer fields if present
+        if self.tokenizer_name:
+            result["tokenizer_name"] = self.tokenizer_name
+        if self.tokenizer_config:
+            result["tokenizer_config"] = self.tokenizer_config
+        if self.embedding_model_name:
+            result["embedding_model_name"] = self.embedding_model_name
+
+        return result
+
+    @staticmethod
+    def from_dict(data: dict, file_path: Path | None = None) -> "ModelMetadata":
+        """Create from dictionary."""
+        return ModelMetadata(
+            model_id=data["model_id"],
+            name=data["name"],
+            version=data["version"],
+            variant=data["variant"],
+            description=data["description"],
+            file_info=FileInfo(
+                filename=data["file_info"]["filename"],
+                size_mb=data["file_info"].get("size_mb", 0.0),
+                checksum=data["file_info"].get("checksum"),
+                onnx_embeddings=data["file_info"].get("onnx_embeddings"),
+            ),
+            performance=PerformanceMetrics(
+                target_latency_ms=data["performance"]["target_latency_ms"],
+                p50_latency_ms=data["performance"].get("p50_latency_ms"),
+                p95_latency_ms=data["performance"].get("p95_latency_ms"),
+                p99_latency_ms=data["performance"].get("p99_latency_ms"),
+                throughput_per_sec=data["performance"].get("throughput_per_sec"),
+                memory_mb=data["performance"].get("memory_mb"),
+            ),
+            accuracy=AccuracyMetrics(
+                binary_f1=data["accuracy"]["binary_f1"] if data.get("accuracy") else None,
+                family_f1=data["accuracy"]["family_f1"] if data.get("accuracy") else None,
+                subfamily_f1=data["accuracy"]["subfamily_f1"] if data.get("accuracy") else None,
+            ) if data.get("accuracy") else None,
+            requirements=Requirements(
+                runtime=ModelRuntime(data["requirements"]["runtime"]),
+                requires_gpu=data["requirements"].get("requires_gpu", False),
+                requires_quantization_support=data["requirements"].get("requires_quantization_support", False),
+            ),
+            status=ModelStatus(data.get("status", "active")),
+            tags=data.get("tags", []),
+            recommended_for=data.get("recommended_for", []),
+            not_recommended_for=data.get("not_recommended_for", []),
+            tokenizer_name=data.get("tokenizer_name"),
+            tokenizer_config=data.get("tokenizer_config"),
+            embedding_model_name=data.get("embedding_model_name"),
+            file_path=file_path,
+        )