swaleha19/agent_tuning_framework

1"""2Domain-Specific Calibration Module for LLMs3 4This module implements calibration techniques for improving uncertainty estimates5across different domains, focusing on temperature scaling and domain adaptation.6"""7 8import numpy as np9import torch10from typing import List, Dict, Any, Union, Optional, Tuple11from scipy.optimize import minimize_scalar12 13class Calibrator:14    """Base class for calibration methods."""15    16    def __init__(self, name: str):17        """18        Initialize the calibrator.19        20        Args:21            name: Name of the calibration method22        """23        self.name = name24        self.is_fitted = False25    26    def fit(self, confidences: List[float], accuracies: List[bool]) -> None:27        """28        Fit the calibrator to the provided data.29        30        Args:31            confidences: List of confidence scores32            accuracies: List of boolean accuracy indicators33        """34        raise NotImplementedError("Subclasses must implement this method")35    36    def calibrate(self, confidences: List[float]) -> List[float]:37        """38        Calibrate the provided confidence scores.39        40        Args:41            confidences: List of confidence scores42            43        Returns:44            Calibrated confidence scores45        """46        raise NotImplementedError("Subclasses must implement this method")47 48 49class TemperatureScaling(Calibrator):50    """Calibration using temperature scaling."""51    52    def __init__(self):53        """Initialize the temperature scaling calibrator."""54        super().__init__("temperature_scaling")55        self.temperature = 1.056    57    def _nll_loss(self, temperature: float, confidences: np.ndarray, accuracies: np.ndarray) -> float:58        """59        Calculate negative log likelihood loss for temperature scaling.60        61        Args:62            temperature: Temperature parameter63            confidences: Array of confidence scores64            accuracies: Array of boolean accuracy indicators65            66        Returns:67            Negative log likelihood loss68        """69        # Apply temperature scaling70        scaled_confidences = np.clip(confidences / temperature, 1e-10, 1.0 - 1e-10)71        72        # Calculate binary cross-entropy loss73        loss = -np.mean(74            accuracies * np.log(scaled_confidences) + 75            (1 - accuracies) * np.log(1 - scaled_confidences)76        )77        78        return loss79    80    def fit(self, confidences: List[float], accuracies: List[bool]) -> None:81        """82        Fit the temperature parameter to the provided data.83        84        Args:85            confidences: List of confidence scores86            accuracies: List of boolean accuracy indicators87        """88        if not confidences or len(confidences) != len(accuracies):89            raise ValueError("Confidences and accuracies must have the same non-zero length")90        91        # Convert to numpy arrays92        conf_array = np.array(confidences)93        acc_array = np.array(accuracies, dtype=float)94        95        # Optimize temperature parameter96        result = minimize_scalar(97            lambda t: self._nll_loss(t, conf_array, acc_array),98            bounds=(0.1, 10.0),99            method='bounded'100        )101        102        self.temperature = result.x103        self.is_fitted = True104        105        print(f"Fitted temperature parameter: {self.temperature:.4f}")106    107    def calibrate(self, confidences: List[float]) -> List[float]:108        """109        Calibrate the provided confidence scores using temperature scaling.110        111        Args:112            confidences: List of confidence scores113            114        Returns:115            Calibrated confidence scores116        """117        if not self.is_fitted:118            raise ValueError("Calibrator must be fitted before calibration")119        120        # Apply temperature scaling121        calibrated = [min(max(conf / self.temperature, 1e-10), 1.0 - 1e-10) for conf in confidences]122        123        return calibrated124 125 126class DomainAdaptiveCalibration(Calibrator):127    """Calibration using domain-adaptive techniques."""128    129    def __init__(self, source_domain: str, target_domain: str):130        """131        Initialize the domain-adaptive calibrator.132        133        Args:134            source_domain: Source domain name135            target_domain: Target domain name136        """137        super().__init__("domain_adaptive_calibration")138        self.source_domain = source_domain139        self.target_domain = target_domain140        self.source_temperature = 1.0141        self.target_temperature = 1.0142        self.domain_shift_factor = 1.0143    144    def fit(145        self, 146        source_confidences: List[float], 147        source_accuracies: List[bool],148        target_confidences: Optional[List[float]] = None,149        target_accuracies: Optional[List[bool]] = None150    ) -> None:151        """152        Fit the domain-adaptive calibrator to the provided data.153        154        Args:155            source_confidences: List of confidence scores from source domain156            source_accuracies: List of boolean accuracy indicators from source domain157            target_confidences: List of confidence scores from target domain (if available)158            target_accuracies: List of boolean accuracy indicators from target domain (if available)159        """160        # Fit source domain temperature161        source_calibrator = TemperatureScaling()162        source_calibrator.fit(source_confidences, source_accuracies)163        self.source_temperature = source_calibrator.temperature164        165        # If target domain data is available, fit target temperature166        if target_confidences and target_accuracies:167            target_calibrator = TemperatureScaling()168            target_calibrator.fit(target_confidences, target_accuracies)169            self.target_temperature = target_calibrator.temperature170            171            # Calculate domain shift factor172            self.domain_shift_factor = self.target_temperature / self.source_temperature173        else:174            # Default domain shift factor based on heuristics175            # This is a simplified approach; in a real system, this would be more sophisticated176            self.domain_shift_factor = 1.2  # Assuming target domain is slightly more uncertain177            self.target_temperature = self.source_temperature * self.domain_shift_factor178        179        self.is_fitted = True180        181        print(f"Fitted source temperature: {self.source_temperature:.4f}")182        print(f"Fitted target temperature: {self.target_temperature:.4f}")183        print(f"Domain shift factor: {self.domain_shift_factor:.4f}")184    185    def calibrate(self, confidences: List[float], domain: str = None) -> List[float]:186        """187        Calibrate the provided confidence scores using domain-adaptive calibration.188        189        Args:190            confidences: List of confidence scores191            domain: Domain of the confidences ('source' or 'target', defaults to target)192            193        Returns:194            Calibrated confidence scores195        """196        if not self.is_fitted:197            raise ValueError("Calibrator must be fitted before calibration")198        199        # Determine which temperature to use200        if domain == "source":201            temperature = self.source_temperature202        else:203            temperature = self.target_temperature204        205        # Apply temperature scaling206        calibrated = [min(max(conf / temperature, 1e-10), 1.0 - 1e-10) for conf in confidences]207        208        return calibrated209 210 211class EnsembleCalibration(Calibrator):212    """Calibration using an ensemble of calibration methods."""213    214    def __init__(self, calibrators: List[Calibrator], weights: Optional[List[float]] = None):215        """216        Initialize the ensemble calibrator.217        218        Args:219            calibrators: List of calibrator instances220            weights: List of weights for each calibrator (None for equal weights)221        """222        super().__init__("ensemble_calibration")223        self.calibrators = calibrators224        225        # Initialize weights226        if weights is None:227            self.weights = [1.0 / len(calibrators)] * len(calibrators)228        else:229            if len(weights) != len(calibrators):230                raise ValueError("Number of weights must match number of calibrators")231            232            # Normalize weights233            total = sum(weights)234            self.weights = [w / total for w in weights]235    236    def fit(self, confidences: List[float], accuracies: List[bool]) -> None:237        """238        Fit all calibrators in the ensemble.239        240        Args:241            confidences: List of confidence scores242            accuracies: List of boolean accuracy indicators243        """244        for calibrator in self.calibrators:245            calibrator.fit(confidences, accuracies)246        247        self.is_fitted = True248    249    def calibrate(self, confidences: List[float]) -> List[float]:250        """251        Calibrate the provided confidence scores using the ensemble.252        253        Args:254            confidences: List of confidence scores255            256        Returns:257            Calibrated confidence scores258        """259        if not self.is_fitted:260            raise ValueError("Calibrator must be fitted before calibration")261        262        # Get calibrated confidences from each calibrator263        all_calibrated = []264        for calibrator in self.calibrators:265            all_calibrated.append(calibrator.calibrate(confidences))266        267        # Combine calibrated confidences using weights268        calibrated = []269        for i in range(len(confidences)):270            weighted_sum = sum(self.weights[j] * all_calibrated[j][i] for j in range(len(self.calibrators)))271            calibrated.append(weighted_sum)272        273        return calibrated274 275 276# Factory function to create calibrators277def create_calibrator(method: str, **kwargs) -> Calibrator:278    """279    Create a calibrator based on the specified method.280    281    Args:282        method: Name of the calibration method283        **kwargs: Additional arguments for the calibrator284        285    Returns:286        Calibrator instance287    """288    if method == "temperature_scaling":289        return TemperatureScaling()290    elif method == "domain_adaptive":291        if "source_domain" not in kwargs or "target_domain" not in kwargs:292            raise ValueError("Domain-adaptive calibration requires source_domain and target_domain")293        return DomainAdaptiveCalibration(kwargs["source_domain"], kwargs["target_domain"])294    elif method == "ensemble":295        if "calibrators" not in kwargs:296            raise ValueError("Ensemble calibration requires a list of calibrators")297        return EnsembleCalibration(kwargs["calibrators"], kwargs.get("weights"))298    else:299        raise ValueError(f"Unsupported calibration method: {method}")300