Initial commit

libs/calibration/multi_model.py

@@ -0,0 +1,85 @@
+"""Multi-model calibrator for handling multiple models/tasks."""
+
+import glob
+import os
+
+import structlog
+
+from .calibrator import ConfidenceCalibrator
+
+logger = structlog.get_logger()
+
+
+class MultiModelCalibrator:
+    """Calibrate confidence scores for multiple models/tasks"""
+
+    def __init__(self) -> None:
+        self.calibrators: dict[str, ConfidenceCalibrator] = {}
+
+    def add_calibrator(self, model_name: str, method: str = "temperature") -> None:
+        """Add calibrator for a specific model"""
+        self.calibrators[model_name] = ConfidenceCalibrator(method)
+        logger.info("Added calibrator", model=model_name, method=method)
+
+    def fit(self, model_name: str, scores: list[float], labels: list[bool]) -> None:
+        """Fit calibrator for specific model"""
+        if model_name not in self.calibrators:
+            self.add_calibrator(model_name)
+
+        self.calibrators[model_name].fit(scores, labels)
+
+    def calibrate(self, model_name: str, scores: list[float]) -> list[float]:
+        """Calibrate scores for specific model"""
+        if model_name not in self.calibrators:
+            logger.warning("No calibrator for model", model=model_name)
+            return scores
+
+        return self.calibrators[model_name].calibrate(scores)
+
+    def save_all(self, directory: str) -> None:
+        """Save all calibrators"""
+        os.makedirs(directory, exist_ok=True)
+
+        for model_name, calibrator in self.calibrators.items():
+            filepath = os.path.join(directory, f"{model_name}_calibrator.pkl")
+            calibrator.save_model(filepath)
+
+    def load_all(self, directory: str) -> None:
+        """Load all calibrators from directory"""
+        pattern = os.path.join(directory, "*_calibrator.pkl")
+        for filepath in glob.glob(pattern):
+            filename = os.path.basename(filepath)
+            model_name = filename.replace("_calibrator.pkl", "")
+
+            calibrator = ConfidenceCalibrator()
+            calibrator.load_model(filepath)
+            self.calibrators[model_name] = calibrator
+
+    def save_models(self, directory: str) -> None:
+        """Save all calibrators (alias for save_all)"""
+        self.save_all(directory)
+
+    def load_models(self, directory: str) -> None:
+        """Load all calibrators from directory (alias for load_all)"""
+        self.load_all(directory)
+
+    def get_model_names(self) -> list[str]:
+        """Get list of model names"""
+        return list(self.calibrators.keys())
+
+    def has_model(self, model_name: str) -> bool:
+        """Check if model exists"""
+        return model_name in self.calibrators
+
+    def is_fitted(self, model_name: str) -> bool:
+        """Check if model is fitted"""
+        if model_name not in self.calibrators:
+            raise ValueError(f"Model '{model_name}' not found")
+        return self.calibrators[model_name].is_fitted
+
+    def remove_calibrator(self, model_name: str) -> None:
+        """Remove calibrator for specific model"""
+        if model_name not in self.calibrators:
+            raise ValueError(f"Model '{model_name}' not found")
+        del self.calibrators[model_name]
+        logger.info("Removed calibrator", model=model_name)