potato/solo_mode/uncertainty/sampling_diversity.py

"""
Sampling Diversity Uncertainty Estimation

This strategy runs the model multiple times at high temperature and measures
the diversity of responses to estimate uncertainty.
"""

import logging
from collections import Counter
from typing import Any, Dict, List, Optional

from .base import UncertaintyEstimator, UncertaintyEstimate

logger = logging.getLogger(__name__)


class SamplingDiversityEstimator(UncertaintyEstimator):
    """
    Estimate uncertainty by sampling multiple responses at high temperature.

    This strategy runs the model N times with high temperature and measures
    how much the predicted labels vary. High diversity in responses indicates
    the model is uncertain about the correct answer.

    Pros:
    - Captures actual model uncertainty in a direct way
    - Works with all model endpoints
    - Provides interpretable diversity metrics

    Cons:
    - Requires N API calls per instance (expensive)
    - Slower than other methods
    - May not work well for deterministic prompts
    """

    def __init__(self, config: Optional[Dict[str, Any]] = None):
        """Initialize the sampling diversity estimator."""
        super().__init__(config)
        self.num_samples = self.config.get('num_samples', 5)
        self.temperature = self.config.get('temperature', 1.0)
        self.normalize_labels = self.config.get('normalize_labels', True)

    def estimate_uncertainty(
        self,
        instance_id: str,
        text: str,
        prompt: str,
        predicted_label: Any,
        endpoint: Any,
        schema_info: Optional[Dict[str, Any]] = None
    ) -> UncertaintyEstimate:
        """
        Estimate uncertainty by sampling multiple responses.

        Makes N API calls at high temperature and measures label diversity.
        """
        try:
            # Get valid labels for normalization
            valid_labels = self._get_valid_labels(schema_info) if schema_info else None

            # Sample multiple responses
            sampled_labels = self._sample_responses(
                prompt, endpoint, valid_labels
            )

            if not sampled_labels:
                return UncertaintyEstimate(
                    uncertainty_score=0.5,
                    confidence_score=0.5,
                    method='sampling_diversity',
                    metadata={'error': 'No valid samples obtained'}
                )

            logger.info(
                f"[SamplingDiversity] {instance_id}: "
                f"{len(sampled_labels)}/{self.num_samples} samples, "
                f"labels={dict(Counter(sampled_labels))}"
            )

            # Calculate diversity metrics
            label_counts = Counter(sampled_labels)
            total_samples = len(sampled_labels)
            unique_labels = len(label_counts)

            # Most common label and its count
            most_common_label, most_common_count = label_counts.most_common(1)[0]

            # Uncertainty = 1 - (most_common_count / total_samples)
            # If all samples agree, uncertainty = 0
            # If samples are evenly distributed, uncertainty approaches 1
            consistency_score = most_common_count / total_samples
            uncertainty_score = 1.0 - consistency_score

            # Alternative: Entropy-based diversity
            entropy = self._calculate_entropy(label_counts, total_samples)
            max_entropy = self._max_entropy(unique_labels)
            normalized_entropy = entropy / max_entropy if max_entropy > 0 else 0

            return UncertaintyEstimate(
                uncertainty_score=uncertainty_score,
                confidence_score=consistency_score,
                method='sampling_diversity',
                metadata={
                    'num_samples': total_samples,
                    'unique_labels': unique_labels,
                    'label_distribution': dict(label_counts),
                    'most_common_label': most_common_label,
                    'most_common_count': most_common_count,
                    'consistency_score': consistency_score,
                    'normalized_entropy': normalized_entropy,
                }
            )

        except Exception as e:
            logger.warning(
                f"Error in sampling diversity estimation for {instance_id}: {e}. "
                "Returning default uncertainty."
            )
            return UncertaintyEstimate(
                uncertainty_score=0.5,
                confidence_score=0.5,
                method='sampling_diversity',
                metadata={'error': str(e)}
            )

    def _sample_responses(
        self,
        prompt: str,
        endpoint: Any,
        valid_labels: Optional[List[str]] = None
    ) -> List[str]:
        """
        Sample multiple responses from the model at high temperature.

        Attempts batched sampling first (vLLM n parameter) for efficiency,
        falls back to sequential calls for other endpoints.

        Args:
            prompt: The labeling prompt
            endpoint: The AI endpoint
            valid_labels: Optional list of valid labels for normalization

        Returns:
            List of sampled labels
        """
        # Try batched sampling first (vLLM/OpenAI support n parameter)
        batched = self._try_batched_sampling(prompt, endpoint, valid_labels)
        if batched is not None:
            return batched

        # Fallback: sequential sampling
        return self._sequential_sampling(prompt, endpoint, valid_labels)

    def _try_batched_sampling(
        self,
        prompt: str,
        endpoint: Any,
        valid_labels: Optional[List[str]] = None
    ) -> Optional[List[str]]:
        """Try to sample N responses in a single API call using the n parameter.

        Works with vLLM and OpenAI-compatible endpoints. Returns None if
        the endpoint doesn't support batched completions.
        """
        # Only attempt for endpoints with base_url (vLLM, OpenAI-compatible)
        base_url = getattr(endpoint, 'base_url', None)
        if not base_url:
            ai_config = getattr(endpoint, 'ai_config', {})
            base_url = ai_config.get('base_url')
        if not base_url:
            return None

        try:
            import requests as req

            headers = {"Content-Type": "application/json"}
            api_key = getattr(endpoint, 'api_key', '')
            if api_key:
                headers["Authorization"] = f"Bearer {api_key}"

            # Get think setting from endpoint config
            ai_config = getattr(endpoint, 'ai_config', {})
            think = ai_config.get('think', False)

            payload = {
                "model": getattr(endpoint, 'model', ''),
                "messages": [{"role": "user", "content": prompt}],
                "max_tokens": getattr(endpoint, 'max_tokens', 200),
                "temperature": self.temperature,
                "stream": False,
                "n": self.num_samples,
                "chat_template_kwargs": {"enable_thinking": think},
            }

            timeout = ai_config.get('timeout', 60)
            response = req.post(
                f"{base_url}/v1/chat/completions",
                headers=headers,
                json=payload,
                timeout=timeout,
            )

            if response.status_code != 200:
                return None

            data = response.json()
            choices = data.get("choices", [])
            if len(choices) < 2:
                return None  # Server didn't support n parameter

            sampled_labels = []
            for choice in choices:
                content = choice.get("message", {}).get("content", "")
                if not content:
                    continue

                # Parse label from response
                label = self._extract_label_from_response(content)
                if label and self.normalize_labels and valid_labels:
                    label = self._normalize_label(label, valid_labels)
                if label:
                    sampled_labels.append(label)

            return sampled_labels if sampled_labels else None

        except Exception as e:
            logger.debug(f"Batched sampling failed, falling back to sequential: {e}")
            return None

    def _sequential_sampling(
        self,
        prompt: str,
        endpoint: Any,
        valid_labels: Optional[List[str]] = None
    ) -> List[str]:
        """Sample responses one at a time (fallback for non-batching endpoints)."""
        sampled_labels = []
        original_temp = getattr(endpoint, 'temperature', 0.1)

        try:
            if hasattr(endpoint, 'temperature'):
                endpoint.temperature = self.temperature

            for i in range(self.num_samples):
                try:
                    from pydantic import BaseModel

                    class LabelResponse(BaseModel):
                        label: str

                    response = endpoint.query(prompt, LabelResponse)
                    label = self._extract_label_from_response(response)

                    if label and self.normalize_labels and valid_labels:
                        label = self._normalize_label(label, valid_labels)
                    if label:
                        sampled_labels.append(label)

                except Exception as e:
                    logger.debug(f"Error in sample {i}: {e}")
                    continue

        finally:
            if hasattr(endpoint, 'temperature'):
                endpoint.temperature = original_temp

        return sampled_labels

    @staticmethod
    def _extract_label_from_response(response) -> Optional[str]:
        """Extract a label string from various response formats."""
        import json as json_mod
        import re

        if isinstance(response, dict):
            return response.get('label', str(response))
        elif isinstance(response, str):
            content = response.strip()
            # Try JSON parse
            try:
                data = json_mod.loads(content)
                if isinstance(data, dict):
                    return data.get('label', '')
            except (json_mod.JSONDecodeError, ValueError):
                pass
            # Try extracting from markdown
            match = re.search(r'"label"\s*:\s*"([^"]+)"', content)
            if match:
                return match.group(1)
            return content
        elif hasattr(response, 'model_dump'):
            return response.model_dump().get('label', str(response))
        elif hasattr(response, 'label'):
            return response.label
        return str(response)

    def _get_valid_labels(self, schema_info: Dict[str, Any]) -> Optional[List[str]]:
        """Extract valid labels from schema info."""
        labels = schema_info.get('labels', [])
        valid_labels = []
        for label in labels:
            if isinstance(label, str):
                valid_labels.append(label)
            elif isinstance(label, dict):
                valid_labels.append(label.get('name', str(label)))
        return valid_labels if valid_labels else None

    def _normalize_label(self, label: Any, valid_labels: List[str]) -> Optional[str]:
        """
        Normalize a label to match one of the valid labels.

        Returns None if no match found. Handles non-string inputs gracefully.
        """
        if label is None:
            return None
        try:
            label_lower = str(label).lower().strip()
        except Exception:
            return None

        for valid in valid_labels:
            if str(valid).lower().strip() == label_lower:
                return valid

        # No match found
        return None

    def _calculate_entropy(
        self,
        label_counts: Counter,
        total: int
    ) -> float:
        """Calculate Shannon entropy of label distribution."""
        import math

        entropy = 0.0
        for count in label_counts.values():
            if count > 0:
                p = count / total
                entropy -= p * math.log2(p)

        return entropy

    def _max_entropy(self, num_classes: int) -> float:
        """Calculate maximum possible entropy for N classes."""
        import math

        if num_classes <= 1:
            return 0.0
        return math.log2(num_classes)

    def supports_endpoint(self, endpoint: Any) -> bool:
        """
        Check if endpoint is supported.

        Sampling diversity works with any text-generation endpoint,
        though it requires setting temperature.
        """
        return True

    def get_config_defaults(self) -> Dict[str, Any]:
        """Get default configuration."""
        return {
            'num_samples': 5,
            'temperature': 1.0,
            'normalize_labels': True,
        }

    def validate_config(self) -> List[str]:
        """Validate configuration."""
        errors = []

        if self.num_samples < 2:
            errors.append("num_samples must be at least 2")

        if self.temperature <= 0:
            errors.append("temperature must be positive")

        return errors