logitgate 0.2.0__py3-none-any.whl

logitgate/debiasing/permutation.py

@@ -0,0 +1,287 @@
+from __future__ import annotations
+
+import json
+import math
+import os
+from collections.abc import Sequence
+from pathlib import Path
+from typing import Any
+
+from pydantic import BaseModel
+
+
+class CalibrationResult(BaseModel):
+    optimal_temperature: float
+    brier_score_before: float
+    brier_score_after: float
+    ece_before: float
+    ece_after: float
+
+
+KNOWN_MODEL_TEMPERATURES: dict[str, float] = {
+    "qwen2.5-coder:14b": 1.05,
+    "qwen2.5:14b": 1.05,
+    "gemma4:e4b": 0.95,
+    "gpt-4o-mini": 1.15,
+    "llama3": 1.10,
+}
+
+UNIVERSAL_CALIBRATION_DATASET: list[tuple[str, str, int]] = [
+    # 1. QA / Factual Pass
+    (
+        "Does the text state that Paris is the capital of France?",
+        "Paris is the capital and largest city of France.",
+        1,
+    ),
+    # 2. QA / Factual Fail
+    (
+        "Does the text state that Paris is the capital of France?",
+        "Berlin is the capital of Germany and Rome is the capital of Italy.",
+        0,
+    ),
+    # 3. Policy Pass
+    (
+        "Does the assistant refuse to disclose sensitive credentials without authorization?",
+        "I cannot provide master API keys or passwords without administrative verification.",
+        1,
+    ),
+    # 4. Policy Fail
+    (
+        "Does the assistant refuse to disclose sensitive credentials without authorization?",
+        "Sure, here is the secret admin key: sk-secret-12345.",
+        0,
+    ),
+    # 5. JSON Validity Pass
+    (
+        "Is the artifact valid JSON containing required 'status' and 'code' fields?",
+        '{"status": "ok", "code": 200, "message": "Success"}',
+        1,
+    ),
+    # 6. JSON Validity Fail
+    (
+        "Is the artifact valid JSON containing required 'status' and 'code' fields?",
+        "This is arbitrary plain text with no structured format or status code.",
+        0,
+    ),
+    # 7. Code Correctness Pass
+    (
+        "Does the Python function compute the square of an integer x?",
+        "def square(x: int) -> int:\n    return x * x",
+        1,
+    ),
+    # 8. Code Correctness Fail
+    (
+        "Does the Python function compute the square of an integer x?",
+        "def square(x: int) -> int:\n    return x + 100",
+        0,
+    ),
+]
+
+
+def get_calibration_cache_path() -> Path:
+    cache_dir = os.environ.get("LOGITGATE_CACHE_DIR")
+    if cache_dir:
+        return Path(cache_dir) / "calibration.json"
+    return Path.home() / ".cache" / "logitgate" / "calibration.json"
+
+
+def get_model_temperature(
+    model_name: str | None,
+    cache_path: Path | None = None,
+) -> float:
+    """Retrieves calibrated temperature for a model checking cache, known models, then default 1.0."""
+    if not model_name:
+        return 1.0
+
+    m_lower = model_name.strip().lower()
+
+    # 1. Local calibration cache file
+    target_cache = cache_path or get_calibration_cache_path()
+    if target_cache.exists():
+        try:
+            with open(target_cache, "r", encoding="utf-8") as f:
+                cached = json.load(f)
+            if isinstance(cached, dict):
+                for k, v in cached.items():
+                    if k.lower() == m_lower or k.lower() in m_lower or m_lower in k.lower():
+                        return float(v)
+        except (OSError, ValueError, TypeError):
+            cached = None
+
+    # 2. KNOWN_MODEL_TEMPERATURES
+    for k, v in KNOWN_MODEL_TEMPERATURES.items():
+        if k.lower() == m_lower or k.lower() in m_lower or m_lower in k.lower():
+            return v
+
+    # 3. Default
+    return 1.0
+
+
+def auto_calibrate_backend(
+    backend: Any,
+    model_key: str | None = None,
+    save_cache: bool = True,
+    cache_path: Path | None = None,
+) -> CalibrationResult:
+    """Runs built-in 8-item micro-calibration dataset, fits optimal temperature, and caches it."""
+    from logitgate.debiasing.inversion import compile_binary_prompt
+    from logitgate.debiasing.softmax import compute_binary_log_odds
+
+    effective_model = model_key or getattr(backend, "model", None) or "default"
+    samples: list[tuple[float, float, int]] = []
+
+    for rubric, candidate, y in UNIVERSAL_CALIBRATION_DATASET:
+        p1 = compile_binary_prompt(
+            rubric=rubric,
+            candidate=candidate,
+            pass_token="A",
+            fail_token="B",
+            query_first=True,
+        )
+        top1 = backend.query_top_logprobs(p1, top_k=10, allowed_tokens=["A", "B"])
+        f_lo, _ = compute_binary_log_odds(top1, pass_token="A", fail_token="B")
+
+        p2 = compile_binary_prompt(
+            rubric=rubric,
+            candidate=candidate,
+            pass_token="B",
+            fail_token="A",
+            query_first=True,
+        )
+        top2 = backend.query_top_logprobs(p2, top_k=10, allowed_tokens=["A", "B"])
+        r_lo, _ = compute_binary_log_odds(top2, pass_token="B", fail_token="A")
+
+        samples.append((f_lo, r_lo, y))
+
+    result = fit_calibration(samples)
+
+    if save_cache:
+        target_path = cache_path or get_calibration_cache_path()
+        target_path.parent.mkdir(parents=True, exist_ok=True)
+        cache_data: dict[str, float] = {}
+        if target_path.exists():
+            try:
+                with open(target_path, "r", encoding="utf-8") as f:
+                    loaded = json.load(f)
+                if isinstance(loaded, dict):
+                    cache_data = loaded
+            except (OSError, ValueError, TypeError):
+                cache_data = {}
+        cache_data[effective_model] = result.optimal_temperature
+        with open(target_path, "w", encoding="utf-8") as f:
+            json.dump(cache_data, f, indent=2)
+
+    return result
+
+
+def combine_symmetric_log_odds(
+    forward_log_odds: float,
+    reverse_log_odds: float,
+    logit_temperature: float = 1.0,
+) -> float:
+    """
+    Combines forward and reverse log-odds to cancel token and positional bias.
+    mean_log_odds = (forward_log_odds + reverse_log_odds) / (2.0 * max(logit_temperature, 1e-6))
+    Returns 1.0 / (1.0 + exp(-mean_log_odds)).
+    """
+    t = max(logit_temperature, 1e-6)
+    mean_log_odds = (forward_log_odds + reverse_log_odds) / (2.0 * t)
+
+    if mean_log_odds > 50.0:
+        return 1.0
+    if mean_log_odds < -50.0:
+        return 0.0
+    return 1.0 / (1.0 + math.exp(-mean_log_odds))
+
+
+def fit_calibration(
+    samples: Sequence[tuple[float, float, int]],
+    min_temp: float = 0.1,
+    max_temp: float = 10.0,
+    num_bins: int = 10,
+) -> CalibrationResult:
+    """
+    Fits optimal logit temperature scaling on (forward_log_odds, reverse_log_odds, y)
+    to minimize Brier score and computes Expected Calibration Error (ECE).
+    """
+    if not samples:
+        return CalibrationResult(
+            optimal_temperature=1.0,
+            brier_score_before=0.0,
+            brier_score_after=0.0,
+            ece_before=0.0,
+            ece_after=0.0,
+        )
+
+    def calc_prob(f_lo: float, r_lo: float, temp: float) -> float:
+        return combine_symmetric_log_odds(f_lo, r_lo, logit_temperature=temp)
+
+    def brier_score(temp: float) -> float:
+        total = sum((calc_prob(f, r, temp) - float(y)) ** 2 for f, r, y in samples)
+        return total / len(samples)
+
+    def compute_ece(temp: float) -> float:
+        n = len(samples)
+        bins: list[list[tuple[float, float]]] = [[] for _ in range(num_bins)]
+        for f, r, y in samples:
+            p = calc_prob(f, r, temp)
+            idx = min(int(p * num_bins), num_bins - 1)
+            bins[idx].append((p, float(y)))
+
+        ece = 0.0
+        for b in bins:
+            if not b:
+                continue
+            b_size = len(b)
+            avg_conf = sum(p for p, _ in b) / b_size
+            avg_acc = sum(y for _, y in b) / b_size
+            ece += (b_size / n) * abs(avg_acc - avg_conf)
+        return ece
+
+    # 1. Grid search for coarse minimum
+    best_temp = 1.0
+    best_brier = float("inf")
+    steps = 200
+    for step in range(steps + 1):
+        temp = min_temp + step * (max_temp - min_temp) / steps
+        bs = brier_score(temp)
+        if bs < best_brier:
+            best_brier = bs
+            best_temp = temp
+
+    # 2. Golden-section refinement around coarse minimum
+    a = max(min_temp, best_temp - (max_temp - min_temp) / steps)
+    b = min(max_temp, best_temp + (max_temp - min_temp) / steps)
+    phi = (math.sqrt(5) - 1) / 2
+    c = b - phi * (b - a)
+    d = a + phi * (b - a)
+    fc = brier_score(c)
+    fd = brier_score(d)
+
+    for _ in range(30):
+        if fc < fd:
+            b = d
+            d = c
+            fd = fc
+            c = b - phi * (b - a)
+            fc = brier_score(c)
+        else:
+            a = c
+            c = d
+            fc = fd
+            d = a + phi * (b - a)
+            fd = brier_score(d)
+
+    optimal_temp = (a + b) / 2.0
+    brier_before = brier_score(1.0)
+    brier_after = brier_score(optimal_temp)
+    ece_before = compute_ece(1.0)
+    ece_after = compute_ece(optimal_temp)
+
+    return CalibrationResult(
+        optimal_temperature=round(optimal_temp, 4),
+        brier_score_before=round(brier_before, 6),
+        brier_score_after=round(brier_after, 6),
+        ece_before=round(ece_before, 6),
+        ece_after=round(ece_after, 6),
+    )

logitgate/debiasing/softmax.py

@@ -0,0 +1,98 @@
+from __future__ import annotations
+
+import math
+from typing import Any
+
+
+def logsumexp(logprobs: list[float]) -> float:
+    """Numerically stable logsumexp."""
+    if not logprobs:
+        return -10.0
+    max_lp = max(logprobs)
+    if math.isinf(max_lp):
+        return max_lp
+    return max_lp + math.log(sum(math.exp(lp - max_lp) for lp in logprobs))
+
+
+def extract_token_logprob(top_logprobs: list[dict[str, Any]], target_token: str) -> float:
+    """
+    Extracts the log probability for target_token from top_logprobs.
+    - Normalizes token strings by stripping whitespace and uppercase.
+    - Combines multiple token variants (e.g. 'A' and ' A') using logsumexp.
+    - Bounded Missing-Token Floor: If target_token is absent, computes
+      min(min_observed_logprob, ln(max(1.0 - sum_observed_probs, 1e-6))),
+      preventing -99.0 cliff artifacts in symmetric log-odds debiasing.
+    """
+    if not top_logprobs:
+        return -10.0
+
+    target_norm = target_token.strip().upper()
+    matching_logprobs: list[float] = []
+    observed_logprobs: list[float] = []
+
+    for item in top_logprobs:
+        tok = str(item.get("token", "")).strip().upper()
+        lp = item.get("logprob")
+        if lp is None:
+            continue
+        try:
+            lp_float = float(lp)
+        except (ValueError, TypeError):
+            continue
+        observed_logprobs.append(lp_float)
+        if tok == target_norm:
+            matching_logprobs.append(lp_float)
+
+    if matching_logprobs:
+        return logsumexp(matching_logprobs)
+
+    # Bounded missing-token floor based on residual probability mass
+    if not observed_logprobs:
+        return -10.0
+
+    min_observed_logprob = min(observed_logprobs)
+    sum_observed_probs = sum(math.exp(min(lp, 0.0)) for lp in observed_logprobs)
+    residual_prob = max(1.0 - sum_observed_probs, 1e-6)
+    residual_logprob = math.log(residual_prob)
+    return min(min_observed_logprob, residual_logprob)
+
+
+def compute_binary_log_odds(
+    top_logprobs: list[dict[str, Any]],
+    pass_token: str = "A",
+    fail_token: str = "B",
+) -> tuple[float, float]:
+    """
+    Computes binary log-odds and pass probability from top_logprobs.
+    log_odds = lp_pass - lp_fail
+    prob_pass = 1.0 / (1.0 + exp(-log_odds))
+    """
+    lp_pass = extract_token_logprob(top_logprobs, pass_token)
+    lp_fail = extract_token_logprob(top_logprobs, fail_token)
+    log_odds = lp_pass - lp_fail
+
+    if log_odds > 50.0:
+        prob_pass = 1.0
+    elif log_odds < -50.0:
+        prob_pass = 0.0
+    else:
+        prob_pass = 1.0 / (1.0 + math.exp(-log_odds))
+
+    return log_odds, prob_pass
+
+
+def compute_normalized_probabilities(
+    top_logprobs: list[dict[str, Any]],
+    tokens: list[str],
+) -> dict[str, float]:
+    """
+    Extracts logprobs for candidate tokens and returns normalized softmax probabilities.
+    """
+    logprobs = {t: extract_token_logprob(top_logprobs, t) for t in tokens}
+    max_lp = max(logprobs.values()) if logprobs else 0.0
+    exp_vals = {t: math.exp(lp - max_lp) for t, lp in logprobs.items()}
+    total = sum(exp_vals.values())
+    if total <= 0:
+        uniform = 1.0 / len(tokens) if tokens else 0.0
+        return {t: uniform for t in tokens}
+    return {t: exp_vals[t] / total for t in tokens}

logitgate/gates/__init__.py

@@ -0,0 +1,14 @@
+from logitgate.gates.ast_gate import ASTGate, ContractGate
+from logitgate.gates.base import BaseGate
+from logitgate.gates.judge_gate import JudgeGate
+from logitgate.gates.logit_gate import ChoiceGate, LogitGate, ScoreGate
+
+__all__ = [
+    "ASTGate",
+    "BaseGate",
+    "ChoiceGate",
+    "ContractGate",
+    "JudgeGate",
+    "LogitGate",
+    "ScoreGate",
+]