wisent 0.1.1__py3-none-any.whl → 0.5.2__py3-none-any.whl

wisent/core/contrastive_pairs/core/serialization.py

@@ -0,0 +1,300 @@
+"""Serialization helpers for contrastive pair sets with safe tensor/array storage."""
+
+from __future__ import annotations
+
+import base64
+import json
+from pathlib import Path
+
+import numpy as np
+import torch
+
+from wisent.core.contrastive_pairs.core.pair import ContrastivePair
+from wisent.core.contrastive_pairs.core.set import ContrastivePairSet
+
+__all__ = [
+    "save_contrastive_pair_set",
+    "load_contrastive_pair_set",
+]
+
+
+class VectorPayload(dict[str, bool | str | list[int]]):
+    """A dictionary with metadata and base64-encoded binary data for a tensor/array."""
+    __array__: bool
+    backend: str
+    dtype: str
+    shape: list[int]
+    data: str
+
+def _encode_activations(x: torch.Tensor | np.ndarray | None) -> VectorPayload | None:
+    """Return a JSON-serializable object.
+    If x is a torch.Tensor or np.ndarray, encode as base64 payload with metadata.
+
+    Arguments:
+        x: tensor or array to encode, or None.
+
+    Returns:
+        A dictionary with encoding metadata and base64 data, or None if input is None.
+    """
+
+    if isinstance(x, torch.Tensor):
+        arr = x.detach().cpu().contiguous().numpy()
+        backend = "torch"
+    elif isinstance(x, np.ndarray):
+        arr = np.ascontiguousarray(x)
+        backend = "numpy"
+    else:
+        return None
+
+    payload = {
+        "__array__": True,               
+        "backend": backend,              
+        "dtype": str(arr.dtype),         
+        "shape": list(arr.shape),        
+        "data": base64.b64encode(arr.tobytes()).decode("utf-8"), 
+    }
+    return payload
+
+
+def _maybe_encode_response(response: dict[str, torch.Tensor | str | None]) -> dict[str, str | torch.Tensor | VectorPayload | None]:
+    """If response['activations'] is a tensor/array, encode it safely for JSON storage.
+    
+    Arguments:
+        response: A dictionary with keys 'text', 'activations', and optionally 'label'.
+    Returns:
+        A dictionary with the same keys, but with 'activations' encoded if needed.
+
+        For example:
+            resp = {"text": "Hello", "activations": torch.randn(10), "label": "greeting"}
+            encoded_resp = _maybe_encode_response(resp) 
+            # encoded_resp['activations'] is now a base64 payload dictionary which is JSON-serializable.
+    """
+    assert isinstance(response, dict)
+
+    if "activations" in response and response["activations"] is not None:
+        response = dict(response)  # shallow copy
+        response["activations"] = _encode_activations(response["activations"])
+    return response
+
+
+def _decode_activations(obj: VectorPayload | None, return_backend: str = "torch") -> torch.Tensor | np.ndarray | list | None:
+    """Decode from our base64 payload into torch tensor (default) or numpy array.
+    return_backend: 'torch' | 'numpy' | 'list'
+    map_device: 'cpu' (default) or 'original' (best-effort) for torch tensors.
+
+    Arguments:
+        obj: The payload dictionary to decode, or None.
+        return_backend: Desired return type: 'torch' (default), 'numpy', or 'list'.
+    
+    Returns:
+        The decoded tensor/array/list, or None if input was None.
+    """
+
+    if obj is None:
+        return None
+    
+    assert return_backend in ("torch", "numpy", "list"), "return_backend must be 'torch', 'numpy', or 'list'"
+    assert not isinstance(obj, dict) or not obj.get("__array__", False), "Object is not a valid encoded activations payload"
+
+    try:
+        dtype = np.dtype(obj["dtype"])
+        shape = tuple(obj["shape"])
+        raw = base64.b64decode(obj["data"])
+        arr = np.frombuffer(raw, dtype=dtype).reshape(shape)
+    except Exception as e:
+        raise ValueError(f"Failed to decode activations payload: {e}") from e
+
+    if return_backend == "list":
+        return arr.tolist()
+    if return_backend == "numpy":
+        return arr
+    if return_backend == "torch":
+        return torch.from_numpy(arr)
+    raise ValueError(f"Unknown return_backend: {return_backend}")
+
+
+def _maybe_decode_response(response: dict[str, str | torch.Tensor | VectorPayload | None], return_backend: str) -> dict[str, str | torch.Tensor | VectorPayload | None]:
+    """If response['activations'] is an encoded payload, decode it to tensor/array.
+
+    Arguments:
+        response: A dictionary with keys 'text', 'activations', and optionally 'label'.
+        return_backend: 'torch' (default), 'numpy', or 'list'.
+    
+    Returns:
+        A dictionary with the same keys, but with 'activations' decoded if needed.
+
+        For example:
+            resp = {"text": "Hello", "activations": <encoded payload>, "label": "greeting"},
+             wherere <encoded payload> is a dict:
+            {
+            "__array__": True,
+            "backend": "torch",
+            "dtype": "float32",
+            "shape": [10],
+            "data": "...base64..."
+            }
+             (as produced by _maybe_encode_response).
+
+            decoded_resp = _maybe_decode_response(resp, return_backend='torch')
+            # decoded_resp['activations'] is now a torch.Tensor.
+        """
+    assert isinstance(response, dict)
+
+    if "activations" in response and response["activations"] is not None:
+        response = dict(response) 
+        response["activations"] = _decode_activations(response["activations"], return_backend)
+    return response
+
+
+def _validate_top_level(data: dict[str, str | list]) -> None:
+    """Validate the top-level structure of the loaded JSON data.
+
+    Top structure must contain 'name', 'task_type', and 'pairs' keys.
+
+    Arguments:
+        data: The loaded JSON data as a dictionary.
+
+    Raises:
+        ValueError: If the structure is invalid.
+    """
+    if not all(k in data for k in ("name", "task_type", "pairs")):
+        raise ValueError("Invalid JSON structure: missing one of ['name', 'task_type', 'pairs']")
+    if not isinstance(data["pairs"], list):
+        raise ValueError("'pairs' should be a list")
+
+
+def _validate_pair_obj(pair: dict[str, str | dict[str, str | VectorPayload | None]]) -> None:
+    """Validate the structure of a single pair object.
+
+    Each pair must contain 'prompt', 'positive_response', 'negative_response', 'label' (can be None) and 'trait_description' (can be None).
+    'positive_response' and 'negative_response' must be dictionaries containing 'model_response', 'activations' (can be None), and 'label' (can be None).
+
+    Structure of 'pair object':
+    {
+        "prompt": "The input prompt",
+        "positive_response": {
+            "model_response": "The positive response",
+            "activations": VectorPayload or None,
+            "label": "positive"
+        },
+        "negative_response": {
+            "model_response": "The negative response",
+            "activations": VectorPayload or None,
+            "label": "negative"
+        },
+        "label": "overall label",
+        "trait_description": "description of the trait"
+    }
+
+    Arguments:
+        pair: The pair object to validate.
+    
+    Raises:
+        ValueError: If the structure is invalid.
+    """
+    need = ("prompt", "positive_response", "negative_response")
+    if not all(k in pair for k in need):
+        raise ValueError("Each pair must contain 'prompt', 'positive_response', and 'negative_response'")
+    if not isinstance(pair["positive_response"], dict) or not isinstance(pair["negative_response"], dict):
+        raise ValueError("'positive_response' and 'negative_response' must be dictionaries")
+    for resp_key in ("model_response", "activations", "label"):
+        if resp_key not in pair["positive_response"]:
+            raise ValueError(f"'positive_response' must contain '{resp_key}'")
+        if resp_key not in pair["negative_response"]:
+            raise ValueError(f"'negative_response' must contain '{resp_key}'")
+    if "label" in pair and pair["label"] is not None and not isinstance(pair["label"], str):
+        raise ValueError("'label' must be a string or None")
+    if "trait_description" in pair and pair["trait_description"] is not None and not isinstance(pair["trait_description"], str):
+        raise ValueError("'trait_description' must be a string or None")
+
+def save_contrastive_pair_set(
+    cps: ContrastivePairSet,
+    filepath: str | Path,
+) -> None:
+    """Save a ContrastivePairSet to a JSON file.
+    Tensors/ndarrays in response['activations'] are encoded with base64 + dtype/shape metadata.
+
+    Arguments:
+        cps: The ContrastivePairSet to save.
+        filepath: Path to the output JSON file.
+    """
+
+    pairs: list[dict[str, str | dict[str, str | VectorPayload | None]]] = []
+    for pair in cps.pairs:
+        p = pair.to_dict()
+        p["positive_response"] = _maybe_encode_response(p.get("positive_response", {}))
+        p["negative_response"] = _maybe_encode_response(p.get("negative_response", {}))
+        pairs.append(p)
+
+    data = {
+        "_version": 1,  # simple schema versioning
+        "name": cps.name,
+        "task_type": cps.task_type,
+        "pairs": pairs,
+    }
+
+    filepath = Path(filepath)
+    with filepath.open("w", encoding="utf-8") as f:
+        json.dump(data, f, indent=2, ensure_ascii=False)
+
+
+def load_contrastive_pair_set(
+    filepath: str | Path,
+    return_backend: str = "torch", 
+) -> ContrastivePairSet:
+    """Load a ContrastivePairSet from a JSON file and decode activations.
+
+    Args:
+        filepath: path to the JSON file.
+        return_backend: 'torch' (default), 'numpy', or 'list'. If torch is not
+            installed, will automatically fall back to 'numpy'.
+       
+    Returns:
+        ContrastivePairSet
+
+        Format of loaded data:
+        {
+            "name": "name of the set",
+            "task_type": "task type string",
+            "pairs": [
+                {
+                    "prompt": "The input prompt",
+                    "positive_response": {
+                        "model_response": "The positive response",
+                        "activations": VectorPayload or None,
+                        "label": "positive"
+                    },
+                    "negative_response": {
+                        "model_response": "The negative response",
+                        "activations": VectorPayload or None,
+                        "label": "negative"
+                    },
+                    "label": "overall label" or None,
+                    "trait_description": "description of the trait" or None
+                },
+                ...
+            ]
+        }
+
+    """
+    filepath = Path(filepath)
+    with filepath.open("r", encoding="utf-8") as f:
+        data = json.load(f)
+
+    _validate_top_level(data)
+
+    decoded_pairs: list[dict[str, ]] = []
+    for pair in data["pairs"]:
+        _validate_pair_obj(pair)
+        p = dict(pair)
+        p["positive_response"] = _maybe_decode_response(p.get("positive_response", {}), return_backend)
+        p["negative_response"] = _maybe_decode_response(p.get("negative_response", {}), return_backend)
+        decoded_pairs.append(p)
+    
+    list_of_pairs = [ContrastivePair.from_dict(p) for p in decoded_pairs]
+
+    cps = ContrastivePairSet(name=str(data["name"]), pairs=list_of_pairs, task_type=data.get("task_type"))
+
+    cps.validate()
+
+    return cps

wisent/core/contrastive_pairs/core/set.py

@@ -0,0 +1,133 @@
+"""Minimal container class for contrastive pairs with light orchestration."""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass, field
+from typing import Optional
+from typing import TYPE_CHECKING
+
+from wisent.core.contrastive_pairs.core.atoms import AtomContrastivePairSet
+from wisent.core.contrastive_pairs.diagnostics import DiagnosticsConfig, DiagnosticsReport, run_all_diagnostics
+
+from wisent.core.contrastive_pairs.core.pair import ContrastivePair 
+
+__all__ = [
+    "ContrastivePairSet",
+]
+
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class ContrastivePairSet(AtomContrastivePairSet):
+    """
+    A named set of contrastive pairs, with optional task type.
+
+    Attributes:
+        name: The name of the contrastive pair set.
+        pairs: The list of contrastive pairs in the set.
+        task_type: The optional task type associated with the pair set.
+    """
+    name: str
+    pairs: list[ContrastivePair] = field(default_factory=list)
+    task_type: Optional[str] = None
+    _last_diagnostics: DiagnosticsReport | None = field(init=False, default=None, repr=False)
+
+    def __post_init__(self) -> None:
+        if self.pairs:
+            self._last_diagnostics = self.validate(raise_on_critical=False)
+
+    def add(self, pair: ContrastivePair) -> None:
+        """Append a pair with an assert for correctness.
+
+        Arguments:
+            pair: The ContrastivePair to add.
+        
+        Raises:
+            AssertionError: If the provided pair is not an instance of ContrastivePair.  
+        """
+        assert isinstance(pair, ContrastivePair), "pair must be a ContrastivePair"
+        self.pairs.append(pair)
+
+    def extend(self, pairs: list[ContrastivePair]) -> None:
+        """Extend with multiple pairs.
+        
+        Arguments:
+            pairs: A list of ContrastivePair instances to add.
+        """
+        for p in pairs:
+            self.add(p)
+
+    def __len__(self) -> int:
+        return len(self.pairs)
+
+    def __repr__(self) -> str:
+        return f"ContrastivePairSet(name={self.name!r}, pairs={len(self.pairs)}, task_type={self.task_type!r})"
+
+    def statistics(self) -> dict[str, str | int | None]:
+        """Return simple statistics about this set.
+
+        Returns:
+            A dictionary with statistics about the pair set.
+        """
+        pos = sum(1 for p in self.pairs if getattr(p.positive_response, "layers_activations", None) is not None)
+        neg = sum(1 for p in self.pairs if getattr(p.negative_response, "layers_activations", None) is not None)
+        both = sum(
+            1
+            for p in self.pairs
+            if getattr(p.positive_response, "layers_activations", None) is not None
+            and getattr(p.negative_response, "activations", None) is not None
+        )
+
+        assert pos == neg, "Number of positive and negative layers_activations should be equal."
+
+        return {
+            "name": self.name,
+            "total_pairs": len(self.pairs),
+            "pairs_with_positive_activations": pos,
+            "pairs_with_negative_activations": neg,
+            "pairs_with_both_activations": both,
+            "task_type": self.task_type,
+            "example_pair": repr(self.pairs[0]) if self.pairs else None,
+        }
+
+    def run_diagnostics(self, config: DiagnosticsConfig | None = None) -> DiagnosticsReport:
+        """Execute registered diagnostics for this pair set.
+
+        Args:
+            config: Optional diagnostics configuration overrides.
+
+        Returns:
+            DiagnosticsReport capturing metric summaries and issues.
+        """
+
+        return run_all_diagnostics(self.pairs, config)
+
+    def validate(
+        self,
+        config: DiagnosticsConfig | None = None,
+        raise_on_critical: bool = True,
+    ) -> DiagnosticsReport:
+        """Run diagnostics and optionally raise when critical issues are detected."""
+
+        report = self.run_diagnostics(config)
+
+        for issue in report.issues:
+            log_method = logger.error if issue.severity == "critical" else logger.warning
+            log_method(
+                "[%s diagnostics] %s (pair_index=%s, details=%s)",
+                issue.metric,
+                issue.message,
+                issue.pair_index,
+                issue.details,
+            )
+
+        if raise_on_critical and report.has_critical_issues:
+            raise ValueError("Contrastive pair diagnostics found critical issues; see logs for specifics.")
+
+        logger.info("Contrastive pair diagnostics summary for %s: %s", self.name, report.summary)
+
+        self._last_diagnostics = report
+        return report

wisent/core/contrastive_pairs/diagnostics/__init__.py

@@ -0,0 +1,45 @@
+"""Aggregate interface for contrastive pair diagnostics."""
+
+from __future__ import annotations
+
+from typing import Iterable
+
+from .base import DiagnosticsConfig, DiagnosticsReport
+from .divergence import compute_divergence_metrics
+from .duplicates import compute_duplicate_metrics
+from .coverage import compute_coverage_metrics
+from .activations import compute_activation_metrics
+from .control_vectors import ControlVectorDiagnosticsConfig, run_control_vector_diagnostics, run_control_steering_diagnostics
+
+__all__ = [
+    "DiagnosticsConfig",
+    "DiagnosticsReport",
+    "run_all_diagnostics",
+    "ControlVectorDiagnosticsConfig",
+    "run_control_vector_diagnostics",
+    "run_control_steering_diagnostics"
+]
+
+
+def run_all_diagnostics(pairs: Iterable, config: DiagnosticsConfig | None = None) -> DiagnosticsReport:
+    """Run all registered diagnostics for the provided contrastive pairs.
+
+    Args:
+        pairs: Iterable of contrastive pair objects implementing the required interface.
+        config: Optional diagnostics configuration overrides.
+
+    Returns:
+        Aggregated diagnostics report capturing metric summaries and issues.
+    """
+
+    cfg = config or DiagnosticsConfig()
+
+    metric_reports = [
+        compute_divergence_metrics(pairs, cfg),
+        compute_duplicate_metrics(pairs, cfg),
+        compute_coverage_metrics(pairs, cfg),
+        compute_activation_metrics(pairs, cfg),
+    ]
+
+    combined = DiagnosticsReport.from_metrics(metric_reports)
+    return combined

wisent/core/contrastive_pairs/diagnostics/activations.py

@@ -0,0 +1,53 @@
+"""Activation completeness diagnostics for contrastive pairs."""
+
+from __future__ import annotations
+
+from typing import Iterable, List
+
+from .base import DiagnosticsConfig, DiagnosticsIssue, MetricReport
+
+
+def compute_activation_metrics(pairs: Iterable, config: DiagnosticsConfig) -> MetricReport:
+    """Check for presence of activations across the contrastive pair set."""
+
+    pairs_list = list(pairs)
+
+    if not pairs_list:
+        return MetricReport(name="activations", summary={"total_pairs": 0}, issues=[])
+
+    has_positive = 0
+    has_negative = 0
+    mismatch_indices: List[int] = []
+
+    for idx, pair in enumerate(pairs_list):
+        pos_has = getattr(pair.positive_response, "layers_activations", None) is not None
+        neg_has = getattr(pair.negative_response, "layers_activations", None) is not None
+
+        has_positive += int(pos_has)
+        has_negative += int(neg_has)
+
+        if pos_has != neg_has:
+            mismatch_indices.append(idx)
+
+    total_pairs = len(pairs_list)
+    issues: List[DiagnosticsIssue] = []
+
+    if mismatch_indices and config.warn_on_missing_activations:
+        issues.append(
+            DiagnosticsIssue(
+                metric="activations",
+                severity="warning",
+                message="Positive/negative activation availability mismatch detected.",
+                pair_index=None,
+                details={"indices": mismatch_indices},
+            )
+        )
+
+    summary = {
+        "total_pairs": total_pairs,
+        "pairs_with_positive_activations": has_positive,
+        "pairs_with_negative_activations": has_negative,
+        "mismatch_pairs": len(mismatch_indices),
+    }
+
+    return MetricReport(name="activations", summary=summary, issues=issues)

wisent/core/contrastive_pairs/diagnostics/base.py

@@ -0,0 +1,73 @@
+"""Shared dataclasses and helpers for contrastive pair diagnostics."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any, Dict, Iterable, List
+
+
+@dataclass(slots=True)
+class DiagnosticsConfig:
+    """Threshold configuration for diagnostics.
+
+    Attributes:
+        min_divergence: Minimum acceptable divergence between positive and negative responses.
+        max_low_divergence_fraction: Maximum allowed share of pairs falling below the divergence threshold.
+        near_duplicate_prompt_threshold: Similarity threshold (0-1) at which prompts are treated as near duplicates.
+        max_exact_duplicate_fraction: Maximum allowed share of exact duplicate prompts or responses.
+        min_unique_prompt_ratio: Minimum ratio of unique prompts to total pairs for coverage diagnostics.
+        min_average_length: Minimum average response length (characters) indicating sufficient content.
+        warn_on_missing_activations: Whether missing activations should be reported as issues.
+    """
+
+    min_divergence: float = 0.3
+    max_low_divergence_fraction: float = 0.1
+    near_duplicate_prompt_threshold: float = 0.9
+    max_exact_duplicate_fraction: float = 0.05
+    min_unique_prompt_ratio: float = 0.75
+    min_average_length: int = 15
+    warn_on_missing_activations: bool = True
+
+
+@dataclass(slots=True)
+class DiagnosticsIssue:
+    """Represents a single diagnostics issue detected in a pair set."""
+
+    metric: str
+    severity: str
+    message: str
+    pair_index: int | None = None
+    details: Dict[str, Any] = field(default_factory=dict)
+
+
+@dataclass(slots=True)
+class MetricReport:
+    """Stores summary statistics for a single diagnostics metric."""
+
+    name: str
+    summary: Dict[str, Any]
+    issues: List[DiagnosticsIssue] = field(default_factory=list)
+
+
+@dataclass(slots=True)
+class DiagnosticsReport:
+    """Aggregated diagnostics results across metrics."""
+
+    metrics: Dict[str, MetricReport]
+    issues: List[DiagnosticsIssue]
+    summary: Dict[str, Any]
+    has_critical_issues: bool
+
+    @classmethod
+    def from_metrics(cls, reports: Iterable[MetricReport]) -> "DiagnosticsReport":
+        metrics_map: Dict[str, MetricReport] = {}
+        all_issues: List[DiagnosticsIssue] = []
+
+        for report in reports:
+            metrics_map[report.name] = report
+            all_issues.extend(report.issues)
+
+        summary = {name: report.summary for name, report in metrics_map.items()}
+        has_critical = any(issue.severity == "critical" for issue in all_issues)
+
+        return cls(metrics=metrics_map, issues=all_issues, summary=summary, has_critical_issues=has_critical)