nexus_os_v2/ckplug_retriever.py

"""
CK-PLUG Integration for NEXUS OS v2
Implements Confidence Gain (CG) as the concrete μ_ret chemical potential.

Paper: arXiv:2503.15888 — Parameters vs. Context: Fine-Grained Control
of Knowledge Reliance in Language Models

Model-specific ε thresholds (from Appendix B):
  LLaMA2-7B: -2   |  LLaMA3-8B: -1
  Mistral-0.3-7B: -1  |  Qwen2.5-7B: -3
  For general use: default ε = -1
"""
import math
import torch
from typing import List, Optional, Dict, Tuple, Callable
from dataclasses import dataclass

@dataclass
class TokenModulation:
    """Result of CK-PLUG token-level modulation."""
    token_id: int
    original_prob: float
    modulated_prob: float
    cg: float                  # Confidence Gain
    H_para: float              # Entropy (query-only)
    H_cont: float              # Entropy (query+retrieval)
    was_modulated: bool        # True if this token was in V_head and CG < threshold
    alpha: float               # Adaptive blending weight


class CKPLUGCoupling:
    """
    Concrete implementation of the retrieval chemical potential μ_ret
    from the NEXUS OS Landau-Ginzburg framework.

    μ_ret(x) = μ_0 * grounding_score(x)
    where grounding_score is derived from CK-PLUG Confidence Gain:
      - CG > 0  → retrieval SUPPORTS parametric knowledge (high grounding)
      - CG < 0  → retrieval CONFLICTS with parametric knowledge (low grounding)
      - |CG|    → magnitude of confidence shift
    """

    def __init__(
        self,
        epsilon: float = -1.0,          # Model-specific detection threshold
        top_k: int = 50,                 # Union top-k for V_head
        mu_0: float = 0.5,              # Base chemical potential (from LG framework)
        device: str = "cpu",
    ):
        self.epsilon = epsilon
        self.top_k = top_k
        self.mu_0 = mu_0
        self.device = device

    @staticmethod
    def entropy(probs: torch.Tensor) -> float:
        """Shannon entropy H = -Σ p_i log₂ p_i."""
        p = probs[probs > 0]
        return float(-(p * torch.log2(p)).sum().item())

    @staticmethod
    def confidence_gain(
        p_query: torch.Tensor,       # p(x | X_q) — parametric only
        p_rag: torch.Tensor,           # p(x | X_r + X_q) — with retrieval
    ) -> Tuple[float, float, float]:
        """
        Returns: (CG, H_para, H_cont)
        CG = H(p(x|X_q)) - H(p(x|X_r+X_q))
        Positive CG  → retrieval supports (reduces entropy)
        Negative CG  → retrieval conflicts (increases entropy)
        """
        H_para = CKPLUGCoupling.entropy(p_query)
        H_cont = CKPLUGCoupling.entropy(p_rag)
        CG = H_para - H_cont
        return CG, H_para, H_cont

    def compute_chemical_potential(
        self,
        p_query: torch.Tensor,
        p_rag: torch.Tensor,
    ) -> float:
        """
        Map CK-PLUG Confidence Gain to Landau-Ginzburg chemical potential μ_ret.

        Logic:
          CG >> 0   → retrieval strongly supports → μ_ret ≈ μ_0 (max grounding)
          CG ≈ 0    → neutral → μ_ret ≈ 0 (no coupling)
          CG << 0   → retrieval conflicts → μ_ret ≈ -μ_0 (adversarial)

        We use a tanh-sigmoid for smooth interpolation:
          μ_ret = μ_0 * tanh(CG / τ) where τ controls transition sharpness.
        """
        CG, _, _ = self.confidence_gain(p_query, p_rag)
        tau = 0.5  # Transition width in nats
        mu_ret = self.mu_0 * math.tanh(CG / tau)
        return mu_ret

    def modulate_token(
        self,
        p_query: torch.Tensor,        # Shape: (vocab_size,)
        p_rag: torch.Tensor,          # Shape: (vocab_size,)
    ) -> Tuple[torch.Tensor, TokenModulation]:
        """
        Apply CK-PLUG token-level modulation (Eq. 7-10 from paper).
        Returns: (modulated_distribution, modulation_metadata)
        """
        CG, H_para, H_cont = self.confidence_gain(p_query, p_rag)

        # Refined detection threshold (Eq. 11 / Appendix B)
        threshold = self.epsilon * abs(H_cont)

        if CG >= threshold:
            # No conflict — pass through RAG distribution unchanged
            return p_rag, TokenModulation(
                token_id=-1, original_prob=0.0, modulated_prob=0.0,
                cg=CG, H_para=H_para, H_cont=H_cont,
                was_modulated=False, alpha=0.0,
            )

        # Conflict detected — apply modulation
        # Eq. 5: Parameter-aware log probability
        q_para = torch.log(p_query + 1e-10)

        # Eq. 6: Context-aware log probability
        q_cont = torch.log((p_rag + 1e-10) / (p_query + 1e-10))

        # Eq. 10: Adaptive alpha
        alpha = H_cont / (H_para + H_cont + 1e-10)
        alpha = float(torch.clamp(torch.tensor(alpha), 0.0, 1.0).item())

        # Build V_head: union of top-k from both distributions
        topk_para = torch.topk(q_para, self.top_k).indices
        topk_cont = torch.topk(q_cont, self.top_k).indices
        V_head = torch.unique(torch.cat([topk_para, topk_cont]))

        # Eq. 8: Modulation function F
        F = torch.full_like(q_para, -float('inf'))
        F[V_head] = alpha * q_para[V_head] + (1.0 - alpha) * q_cont[V_head]

        # Softmax to get modulated distribution
        p_mod = torch.softmax(F, dim=-1)

        # Find most changed token for metadata
        diff = torch.abs(p_rag - p_mod)
        changed_id = int(torch.argmax(diff).item())

        modulation = TokenModulation(
            token_id=changed_id,
            original_prob=float(p_rag[changed_id].item()),
            modulated_prob=float(p_mod[changed_id].item()),
            cg=CG, H_para=H_para, H_cont=H_cont,
            was_modulated=True, alpha=alpha,
        )

        return p_mod, modulation

    def batch_modulate(
        self,
        p_queries: List[torch.Tensor],  # List of (vocab_size,) tensors
        p_rags: List[torch.Tensor],     # Same length
    ) -> List[Tuple[torch.Tensor, TokenModulation]]:
        """Apply CK-PLUG to a batch of token positions."""
        return [self.modulate_token(pq, pr) for pq, pr in zip(p_queries, p_rags)]

    def get_grounding_field(self, p_query: torch.Tensor, p_rag: torch.Tensor) -> float:
        """
        Return the scalar μ_ret value for insertion into Landau-Ginzburg functional.
        This is the key bridge between CK-PLUG (empirical) and NEXUS OS physics.
        """
        return self.compute_chemical_potential(p_query, p_rag)


# Model-specific epsilon presets (from CK-PLUG Appendix B)
CKPLUG_PRESETS = {
    "llama2": -2.0,
    "llama3": -1.0,
    "mistral": -1.0,
    "qwen2.5": -3.0,
    "granite": -1.5,     # Estimated from paper patterns
    "gemma": -1.0,       # Estimated
    "deepseek": -2.0,    # Estimated (large MoE, conservative)
    "default": -1.0,
}


def get_preset_epsilon(model_family: str) -> float:
    """Get recommended epsilon for a model family."""
    key = model_family.lower()
    for k, v in CKPLUG_PRESETS.items():
        if k in key:
            return v
    return CKPLUG_PRESETS["default"]