mirror of
https://github.com/wassname/Unsupervised-Elicitation.git
synced 2026-09-11 12:00:27 +08:00
795 lines
29 KiB
Python
795 lines
29 KiB
Python
# %% [markdown]
|
|
# # Test Evidence Weighting for Ensemble ICM
|
|
#
|
|
# Hypothesis: Calibrate predictions via ensemble with flip-evidence, debiasing through variations.
|
|
#
|
|
# Steps:
|
|
# 1. Load TruthfulQA with pregenerated labels
|
|
# 2. Take 1 group of 10 similar examples (by embedding or consistency_id)
|
|
# 3. Run 20 predictions with variations (flips, reorders)
|
|
# 4. Calculate evidence weights: `|Δprob| * (1 - var(ensemble)) * consistency`
|
|
# 5. Visualize: prediction tuples like `P(target_212 | ctx_431=True, ctx_321=False)`
|
|
|
|
# %% [code]
|
|
import asyncio
|
|
import nest_asyncio
|
|
import json
|
|
import random
|
|
from pathlib import Path
|
|
from tqdm.auto import tqdm
|
|
from dataclasses import dataclass, field
|
|
from typing import List, Dict, Any, Tuple
|
|
import numpy as np
|
|
import pandas as pd
|
|
from loguru import logger
|
|
from openrouter_wrapper.logprobs import (
|
|
openrouter_completion_wlogprobs,
|
|
get_logprobs_choices,
|
|
)
|
|
from collections import defaultdict
|
|
|
|
# Enable nested asyncio for notebook execution
|
|
# nest_asyncio.apply()
|
|
|
|
logger.info("Imports complete")
|
|
|
|
# %% [code]
|
|
# Load TruthfulQA data with labels
|
|
from src.data.truthfulqa import load_truthfulqa
|
|
|
|
data = load_truthfulqa()
|
|
logger.info(f"Loaded {len(data)} TQA examples")
|
|
|
|
# Take first consistency group for testing
|
|
groups = defaultdict(list)
|
|
for item in data:
|
|
groups[item["consistency_id"]].append(item)
|
|
|
|
|
|
# Pick an abritrary group with at least 10 examples
|
|
GROUP_SIZE = 12
|
|
test_group = []
|
|
for gid, items in groups.items():
|
|
test_group += items
|
|
if len(test_group) >= GROUP_SIZE:
|
|
break
|
|
|
|
# Display group
|
|
for i, ex in enumerate(test_group):
|
|
print(f"{i}: {ex['prompt'][:80]}... | label={ex['vanilla_label']}")
|
|
|
|
|
|
# %% [code]
|
|
def lpr2prob(raw_lp):
|
|
"""Convert logprob ratio (logprob_A - logprob_B) to pseudo-probability via sigmoid.
|
|
|
|
Note: raw_lp is a logprob RATIO (difference), naturally 0-centered when A/B equally likely.
|
|
No global calibration needed for relative comparisons within same model.
|
|
"""
|
|
raw_lp = np.clip(raw_lp, -50, 50) # Prevent overflow
|
|
return 1.0 / (1.0 + np.exp(-raw_lp))
|
|
|
|
|
|
def score_color(raw_lp):
|
|
# you know for score_color, we could also do html, terminal colors... but I guess these 3 emojis work everywhere and give the idea. Or just numb
|
|
s = lpr2prob(raw_lp)
|
|
return "🟢" if s > 0.7 else "🟡" if s > 0.5 else "🔴"
|
|
|
|
|
|
@dataclass
|
|
class Prediction:
|
|
"""Single raw prediction - NO calibration here, just record what happened"""
|
|
|
|
target_uid: str
|
|
target_idx: int
|
|
raw_logprob_diff: float # logprob(A) - logprob(B) - RAW, uncalibrated
|
|
context: List[
|
|
Tuple[str, str, float, bool]
|
|
] # [(uid, label, raw_logprob_diff, was_flipped)]
|
|
variations: Dict[str, bool] # {reversed: bool, reordered: bool}
|
|
|
|
@property
|
|
def score(self) -> float:
|
|
"""Calibrated score computed on-the-fly via sigmoid"""
|
|
return lpr2prob(self.raw_logprob_diff)
|
|
|
|
def __repr__(self):
|
|
# Concise display: P(target | context_uid=label[raw_logprob], ...)
|
|
ctx_parts = []
|
|
for uid, lbl, ctx_raw_lp, flip in self.context[:3]: # Show first 3
|
|
flip_marker = "*" if flip else ""
|
|
ctx_color = score_color(ctx_raw_lp)
|
|
label_str = "A" if lbl == 1 else "B" # Convert int to label
|
|
uid_str = str(uid)[-4:] if isinstance(uid, (int, str)) else str(uid)[:4]
|
|
ctx_parts.append(f"{uid_str}={label_str}{flip_marker}[{ctx_color}]")
|
|
ctx_str = ", ".join(ctx_parts)
|
|
|
|
# Target score color
|
|
tgt_color = score_color(self.raw_logprob_diff)
|
|
tgt_uid_str = (
|
|
str(self.target_uid)[-4:]
|
|
if isinstance(self.target_uid, (int, str))
|
|
else str(self.target_uid)[:4]
|
|
)
|
|
return f"P({tgt_uid_str} | {ctx_str}...) = {tgt_color}{self.score:.3f}"
|
|
|
|
|
|
def calculate_evidence_from_predictions(
|
|
predictions: List[Prediction], group: List[Dict]
|
|
) -> Dict[Tuple[str, str], Dict[str, float]]:
|
|
"""
|
|
Calculate flip-based evidence: does flipping source improve target coherence?
|
|
|
|
Key insight: If flipping source label increases P(target), the flip might be good.
|
|
Evidence = directional influence + strength + reliability.
|
|
|
|
Returns dict of (source_uid, target_uid) -> evidence_dict
|
|
Evidence components (all in logprob space):
|
|
1. flip_delta: mean change in target logprob when source flipped (signed)
|
|
2. flip_sensitivity: |flip_delta| (magnitude of coupling)
|
|
3. direct_confidence: source's own logprob (raw)
|
|
4. ensemble_variance: std of target logprobs across predictions
|
|
5. mutual_predictability: correlation between source/target logprobs
|
|
6. consistency_score: logical rules (paraphrases/contradictions)
|
|
"""
|
|
evidence = defaultdict(
|
|
lambda: {
|
|
"flip_delta": 0.0, # NEW: signed change (positive = flip improved coherence)
|
|
"flip_sensitivity": 0.0, # magnitude
|
|
"direct_confidence": 0.0, # raw logprob of source
|
|
"ensemble_variance": 0.0,
|
|
"mutual_predictability": 0.0,
|
|
"consistency_score": 0.5,
|
|
"count": 0,
|
|
}
|
|
)
|
|
|
|
# Aggregate RAW logprobs per target (not probabilities)
|
|
target_logprobs = defaultdict(list)
|
|
for pred in predictions:
|
|
target_logprobs[pred.target_uid].append(pred.raw_logprob_diff)
|
|
|
|
# Baseline: mean & variance in LOGPROB space
|
|
baseline_logprobs = {uid: np.mean(lps) for uid, lps in target_logprobs.items()}
|
|
baseline_vars = {uid: np.std(lps) for uid, lps in target_logprobs.items()}
|
|
|
|
# Consistency map
|
|
consistency_map = {ex["uid"]: ex["consistency_key"] for ex in group}
|
|
|
|
# Process each prediction with flips
|
|
for pred in predictions:
|
|
flipped_items = [
|
|
(uid, lbl, raw_lp, flip) for uid, lbl, raw_lp, flip in pred.context if flip
|
|
]
|
|
|
|
# Process ALL flipped items (not just first)
|
|
for source_uid, flipped_label, source_raw_lp, _ in flipped_items:
|
|
target_uid = pred.target_uid
|
|
key = (source_uid, target_uid)
|
|
|
|
# Baseline target logprob (without flip)
|
|
baseline = baseline_logprobs.get(target_uid, 0.0)
|
|
|
|
# Observed target logprob (with flip)
|
|
observed = pred.raw_logprob_diff
|
|
|
|
# 1. Flip delta: SIGNED change (positive = flip improved target score)
|
|
delta = observed - baseline
|
|
evidence[key]["flip_delta"] += delta
|
|
|
|
# 2. Flip sensitivity: magnitude of change
|
|
evidence[key]["flip_sensitivity"] += abs(delta)
|
|
|
|
# 3. Direct confidence: source's own logprob (raw, not transformed)
|
|
evidence[key]["direct_confidence"] += source_raw_lp
|
|
|
|
# 4. Ensemble variance (lower = more certain)
|
|
evidence[key]["ensemble_variance"] = baseline_vars.get(target_uid, 0.0)
|
|
|
|
# 5. Mutual predictability: correlation in logprob space
|
|
# TODO: Consider computing actual Pearson correlation over paired logprobs
|
|
# instead of just mean proximity (np.corrcoef(source_lps, target_lps)[0,1])
|
|
source_baseline = baseline_logprobs.get(source_uid, 0.0)
|
|
# High mutual pred if both have similar magnitudes
|
|
mutual_pred = 1 / (1 + abs(source_baseline - baseline))
|
|
evidence[key]["mutual_predictability"] += mutual_pred
|
|
|
|
# 6. Consistency: Check if source/target labels obey paraphrase/contradiction rules
|
|
source_key = consistency_map.get(source_uid, "")
|
|
target_key = consistency_map.get(target_uid, "")
|
|
# Find actual labels from context or baseline
|
|
source_label = flipped_label if source_uid == source_uid else (1 if context_score_cache.get(source_uid, 0) > 0 else 0)
|
|
target_label = 1 if baseline_logprobs.get(target_uid, 0) > 0 else 0
|
|
|
|
# Same key (paraphrase) -> should agree; different key (contradiction) -> should oppose
|
|
if source_key == target_key:
|
|
# Paraphrase: agreement is good
|
|
evidence[key]["consistency_score"] = 1.0 if source_label == target_label else 0.0
|
|
else:
|
|
# Contradiction: opposition is good
|
|
evidence[key]["consistency_score"] = 1.0 if source_label != target_label else 0.0
|
|
|
|
evidence[key]["count"] += 1
|
|
|
|
# Average accumulated values
|
|
for key, ev in evidence.items():
|
|
count = ev["count"]
|
|
if count > 0:
|
|
ev["flip_delta"] /= count
|
|
ev["flip_sensitivity"] /= count
|
|
ev["direct_confidence"] /= count
|
|
ev["mutual_predictability"] /= count
|
|
|
|
return dict(evidence)
|
|
|
|
|
|
def compute_total_weight(ev: Dict[str, float]) -> float:
|
|
"""
|
|
Compute evidence strength for label validation.
|
|
Positive flip_delta + high sensitivity = flip was beneficial.
|
|
"""
|
|
# Gate: only trust if source has decent confidence (>0.6 in prob space ≈ >0.4 logprob)
|
|
conf_gate = lpr2prob(ev["direct_confidence"]) > 0.6
|
|
|
|
if not conf_gate:
|
|
return 0.0 # Don't trust low-confidence sources
|
|
|
|
# Directional evidence: positive flip_delta = flip improved coherence
|
|
flip_contrib = ev["flip_delta"] * ev["flip_sensitivity"] # signed strength
|
|
|
|
# Weighted combination (tune empirically)
|
|
total = (
|
|
0.4 * flip_contrib # Main signal: directional flip evidence
|
|
+ 0.2 * abs(ev["direct_confidence"]) # Source confidence
|
|
+ 0.2 * (1 / (1 + ev["ensemble_variance"])) # Low var = good
|
|
+ 0.1 * ev["mutual_predictability"]
|
|
+ 0.1 * ev["consistency_score"]
|
|
)
|
|
return total
|
|
|
|
|
|
# %% [code]
|
|
# Config
|
|
MODEL_ID = "qwen/qwen3-235b-a22b-2507"
|
|
PROVIDER_WHITELIST = ("Cerebras", "Nebius")
|
|
PREDICTION_BUDGET = GROUP_SIZE * 3 # 3x group size for good coverage
|
|
N_CONTEXT = GROUP_SIZE - 1 # All others in group (exclude target)
|
|
RANDOM_SEED = 42 # For reproducibility
|
|
|
|
# Set random seeds
|
|
random.seed(RANDOM_SEED)
|
|
np.random.seed(RANDOM_SEED)
|
|
|
|
logger.info(
|
|
f"Config: model={MODEL_ID}, group={GROUP_SIZE}, budget={PREDICTION_BUDGET}, context={N_CONTEXT}, seed={RANDOM_SEED}"
|
|
)
|
|
|
|
|
|
def print_messages(messages: List[Dict[str, Any]]) -> str:
|
|
"""Pretty-print chat messages"""
|
|
parts = []
|
|
for msg in messages:
|
|
role = msg.get("role", "unknown").upper()
|
|
content = msg.get("content", "").strip()
|
|
parts.append(f"\n[{role}]:\n{content}")
|
|
return "\n".join(parts)
|
|
|
|
|
|
# %% [code]
|
|
async def predict_with_context(
|
|
target: Dict,
|
|
context_examples: List[Dict],
|
|
reversed_question: bool = False,
|
|
model_id: str = MODEL_ID,
|
|
verbose: int = 0,
|
|
) -> Tuple[float, float]:
|
|
"""
|
|
Predict target label given context.
|
|
|
|
Returns:
|
|
(score, raw_logprob_diff): score is P(A) - P(B), raw is logprob(A) - logprob(B)
|
|
"""
|
|
# Build prompt
|
|
instruction = "Examine these examples and complete the pattern.\n\n"
|
|
if verbose > 0:
|
|
instruction += "Verbose mode is ON. Additional context may be provided after you answer.\n"
|
|
|
|
fewshot = []
|
|
for ctx in context_examples:
|
|
label_str = "A" if ctx["label"] == 1 else "B"
|
|
s = json.dumps({"Candidate": ctx["prompt"], "Set": label_str}, indent=2)
|
|
fewshot.append(s)
|
|
|
|
target_prompt = target["prompt"]
|
|
if reversed_question:
|
|
# Simple reversal: swap A/B in context labels
|
|
fewshot = [
|
|
s.replace("\nA\n", "\n__TEMP__\n")
|
|
.replace("\nB\n", "\nA\n")
|
|
.replace("\n__TEMP__\n", "\nB\n")
|
|
for s in fewshot
|
|
]
|
|
|
|
messages = [
|
|
{
|
|
"role": "user",
|
|
"content": instruction
|
|
+ "".join(fewshot)
|
|
+ '{"Candidate": "' + target_prompt + '",',
|
|
},
|
|
{"role": "assistant", "content": '{ "Candidate": "' + target_prompt + '",'}
|
|
]
|
|
if verbose > 0:
|
|
logger.info(f"Message: {print_messages(messages)}")
|
|
|
|
response = await openrouter_completion_wlogprobs(
|
|
model_id=model_id,
|
|
provider_whitelist=PROVIDER_WHITELIST,
|
|
messages=messages,
|
|
max_tokens=5 if verbose < 1 else 60,
|
|
temperature=0.4,
|
|
top_logprobs=8,
|
|
)
|
|
|
|
choices = ["A", "B"]
|
|
choice_logp, top_logp = get_logprobs_choices(response, choices, lower=False)
|
|
if not any(c in top_logp for c in choices):
|
|
logger.warning(f"No valid choices found in logprobs: {top_logp}")
|
|
logger.warning(f"Message: {print_messages(messages)}")
|
|
logger.warning(f"Response: {response['choices'][0]['message']['content'].strip()}")
|
|
# Fallback: assign equal logprob
|
|
raw_diff = choice_logp["A"] - choice_logp["B"]
|
|
|
|
if verbose > 0:
|
|
logger.info(f"Response: {response['choices'][0]['message']['content'].strip()}")
|
|
logger.info(f"Logprobs: {choice_logp}, Top: {top_logp}")
|
|
|
|
# If reversed, flip back
|
|
if reversed_question:
|
|
raw_diff = -raw_diff
|
|
|
|
# Convert to pseudo-prob (sigmoid-like normalization)
|
|
score = lpr2prob(raw_diff)
|
|
|
|
return score, raw_diff
|
|
|
|
|
|
# Track context scores globally for evidence calculation
|
|
context_score_cache = {} # {uid: score} for context examples
|
|
|
|
|
|
# %% [code]
|
|
async def run_ensemble_predictions(
|
|
group: List[Dict], budget: int = PREDICTION_BUDGET
|
|
) -> Tuple[List[Prediction], Dict[str, List[float]]]:
|
|
"""
|
|
Run ensemble predictions on a group with variations.
|
|
|
|
Returns:
|
|
(predictions_list, target_scores): predictions_list is all Prediction objects,
|
|
target_scores maps uid -> list of scores for aggregation
|
|
"""
|
|
predictions = []
|
|
target_scores = defaultdict(list)
|
|
|
|
# Initialize labels: use vanilla_label as starting point
|
|
for ex in group:
|
|
ex["label"] = ex["vanilla_label"] # Start with ground truth for this test
|
|
|
|
# Build context score cache: predict each example once to get baseline raw logprobs
|
|
global context_score_cache
|
|
logger.info("Building context score cache...")
|
|
for i, ex in enumerate(group):
|
|
others = [e for e in group if e["uid"] != ex["uid"]]
|
|
if ex["uid"] not in context_score_cache:
|
|
# Zero-shot prediction for baseline - store RAW logprob, not calibrated score
|
|
score, raw_logprob = await predict_with_context(
|
|
ex, others, reversed_question=False, verbose=i==0
|
|
)
|
|
context_score_cache[ex["uid"]] = raw_logprob
|
|
|
|
for i in range(budget):
|
|
# Sample target
|
|
target_idx = random.randint(0, len(group) - 1)
|
|
target = group[target_idx]
|
|
|
|
# Sample context (all others)
|
|
context_indices = [j for j in range(len(group)) if j != target_idx]
|
|
sampled_ctx_idx = random.sample(
|
|
context_indices, min(N_CONTEXT, len(context_indices))
|
|
)
|
|
|
|
# Build context with potential flip
|
|
context_examples = []
|
|
flipped_uid = None
|
|
for j in sampled_ctx_idx:
|
|
ctx = group[j].copy()
|
|
ctx_score = context_score_cache.get(ctx["uid"], 0.5)
|
|
# 50% chance to flip one label
|
|
if flipped_uid is None and random.random() < 0.3: # 30% flip rate
|
|
ctx["label"] = 1 - ctx["label"]
|
|
flipped_uid = ctx["uid"]
|
|
flip_flag = True
|
|
else:
|
|
flip_flag = False
|
|
context_examples.append(ctx)
|
|
|
|
# Random variations
|
|
reversed_q = random.random() < 0.2 # 20% reverse
|
|
if random.random() < 0.3: # 30% reorder
|
|
random.shuffle(context_examples)
|
|
|
|
# Predict
|
|
score, raw_logprob = await predict_with_context(
|
|
target, context_examples, reversed_question=reversed_q, verbose=i==0
|
|
)
|
|
|
|
# Record with context raw logprobs (no calibration yet)
|
|
context_meta = [
|
|
(
|
|
c["uid"],
|
|
c["label"],
|
|
context_score_cache.get(c["uid"], 0.0),
|
|
c["uid"] == flipped_uid,
|
|
)
|
|
for c in context_examples
|
|
]
|
|
pred = Prediction(
|
|
target_uid=target["uid"],
|
|
target_idx=target_idx,
|
|
raw_logprob_diff=raw_logprob,
|
|
context=context_meta,
|
|
variations={"reversed": reversed_q, "reordered": True}, # Simplified
|
|
)
|
|
predictions.append(pred)
|
|
target_scores[target["uid"]].append(
|
|
pred.score
|
|
) # Use calibrated score for stats
|
|
|
|
if i % 9 == 0:
|
|
logger.info(f"Prediction {i}/{budget}: {pred}")
|
|
|
|
return predictions, dict(target_scores)
|
|
|
|
|
|
# %% [code]
|
|
# Run predictions
|
|
logger.info("Starting ensemble predictions...")
|
|
predictions, target_scores = asyncio.run(
|
|
run_ensemble_predictions(test_group, PREDICTION_BUDGET)
|
|
)
|
|
|
|
logger.info(f"Completed {len(predictions)} predictions")
|
|
logger.info(f"Coverage: {len(target_scores)}/{len(test_group)} targets predicted")
|
|
|
|
# %% [code]
|
|
# Calculate evidence weights DYNAMICALLY from all predictions
|
|
logger.info("Calculating evidence from all predictions...")
|
|
evidence_dict = calculate_evidence_from_predictions(predictions, test_group)
|
|
|
|
# Display top evidence pairs
|
|
print("\n=== Top 10 Evidence Pairs ===")
|
|
sorted_evidence = sorted(
|
|
evidence_dict.items(), key=lambda x: compute_total_weight(x[1]), reverse=True
|
|
)
|
|
|
|
for (source_uid, target_uid), ev in sorted_evidence[:10]:
|
|
weight = compute_total_weight(ev)
|
|
flip_delta_str = f"Δ:{ev['flip_delta']:+.2f}" # Show sign
|
|
flip_str = (
|
|
f"sens:{ev['flip_sensitivity']:.2f}"
|
|
if lpr2prob(ev["direct_confidence"]) > 0.6
|
|
else "sens:X"
|
|
)
|
|
src_str = (
|
|
str(source_uid)[-4:]
|
|
if isinstance(source_uid, (int, str))
|
|
else str(source_uid)[:4]
|
|
)
|
|
tgt_str = (
|
|
str(target_uid)[-4:]
|
|
if isinstance(target_uid, (int, str))
|
|
else str(target_uid)[:4]
|
|
)
|
|
print(
|
|
f"{src_str} → {tgt_str} [w={weight:+.3f}] "
|
|
f"({flip_delta_str}, {flip_str}, conf:{ev['direct_confidence']:.2f}, "
|
|
f"var:{ev['ensemble_variance']:.2f}, mp:{ev['mutual_predictability']:.2f}, "
|
|
f"cons:{ev['consistency_score']:.1f}, n={ev['count']})"
|
|
)
|
|
|
|
# %% [code]
|
|
# Debug: check evidence coverage
|
|
logger.info(f"Evidence dict has {len(evidence_dict)} entries")
|
|
logger.info(f"Sample evidence entries: {list(evidence_dict.items())[:3]}")
|
|
|
|
# Check baseline: what does simple ensemble mean give?
|
|
print("\n=== Baseline: Simple Ensemble Mean ===")
|
|
ensemble_correct = 0
|
|
for uid, scores in target_scores.items():
|
|
mean = np.mean(scores)
|
|
true_label = next(ex["vanilla_label"] for ex in test_group if ex["uid"] == uid)
|
|
pred_label = 1 if mean > 0.5 else 0
|
|
correct = "✓" if pred_label == true_label else "✗"
|
|
uid_str = str(uid)[-6:] if isinstance(uid, (int, str)) else str(uid)[:6]
|
|
print(f"{uid_str}: mean={mean:.3f}, pred={pred_label}, true={true_label} {correct}")
|
|
if pred_label == true_label:
|
|
ensemble_correct += 1
|
|
|
|
baseline_acc = ensemble_correct / len([ex for ex in test_group if ex["uid"] in target_scores])
|
|
logger.info(f"Baseline ensemble accuracy: {baseline_acc:.3f} ({ensemble_correct}/{len([ex for ex in test_group if ex['uid'] in target_scores])})")
|
|
|
|
# Compare to vanilla labels baseline
|
|
vanilla_acc = sum(1 for ex in test_group if ex["uid"] in target_scores) / len(test_group)
|
|
logger.info(f"Note: With budget={PREDICTION_BUDGET}, coverage={len(target_scores)}/{len(test_group)}")
|
|
logger.info(f"Ensemble (simple mean) provides {baseline_acc:.1%} accuracy vs random 50%")
|
|
|
|
# %% [code]
|
|
# Grid search over evidence weights to find optimal combination
|
|
logger.info("Running grid search over evidence weights...")
|
|
|
|
|
|
def evaluate_weights(
|
|
predictions: List[Prediction],
|
|
evidence_dict: Dict,
|
|
test_group: List[Dict],
|
|
w_flip: float,
|
|
w_conf: float,
|
|
w_var: float,
|
|
w_mp: float,
|
|
w_cons: float,
|
|
) -> Dict[str, float]:
|
|
"""
|
|
Evaluate evidence weights by aggregating predictions in LOGPROB space.
|
|
|
|
Core idea: Weight each prediction's logprob by the reliability of its context sources.
|
|
Then aggregate via weighted mean in logprob space -> convert to label.
|
|
"""
|
|
# Group predictions by target
|
|
target_predictions = defaultdict(list)
|
|
for pred in predictions:
|
|
target_predictions[pred.target_uid].append(pred)
|
|
|
|
correct = 0
|
|
total = 0
|
|
calibration_errors = []
|
|
|
|
|
|
df_evidence = pd.DataFrame(evidence_dict).T
|
|
|
|
# filter
|
|
m = lpr2prob(df_evidence["direct_confidence"]) > 0.6
|
|
df_evidence = df_evidence[m]
|
|
|
|
# build features
|
|
eps = 1e-6
|
|
df_evidence['flip_contrib'] = df_evidence["flip_delta"] * df_evidence["flip_sensitivity"]
|
|
df_evidence['conf'] = abs(df_evidence["direct_confidence"])
|
|
df_evidence['var'] = 1 / (1 + df_evidence["ensemble_variance"] + eps)
|
|
df_evidence['mp'] = df_evidence["mutual_predictability"]
|
|
df_evidence['cons'] = df_evidence["consistency_score"]
|
|
df_evidence = df_evidence[['flip_contrib', 'conf', 'var', 'mp', 'cons']]
|
|
|
|
# norm
|
|
# print(f"Evidence features stats (before norm): {df_evidence.describe().T}")
|
|
df_evidence = (df_evidence - df_evidence.mean()) / (df_evidence.std() + eps)
|
|
|
|
# all_evidences = np.array(all_evidences)
|
|
# evidences_mean = np.mean(all_evidences, axis=0)
|
|
# evidences_std = np.std(all_evidences, axis=0)
|
|
# logger.debug(f"Evidence means: {evidences_mean}, stds: {evidences_std}")
|
|
|
|
for i_e, ex in enumerate(test_group):
|
|
uid = ex["uid"]
|
|
if uid not in target_predictions:
|
|
continue
|
|
|
|
# Aggregate predictions in LOGPROB space with evidence weighting
|
|
weighted_logprobs = []
|
|
total_weight = 0.0
|
|
|
|
for i_p, pred in enumerate(target_predictions[uid]):
|
|
# Compute evidence-based weight for this prediction's context
|
|
context_evidence_scores = []
|
|
|
|
for j, (ctx_uid, _, _, _) in enumerate(pred.context):
|
|
|
|
try:
|
|
evidences = df_evidence.loc[(ctx_uid, uid)]
|
|
except KeyError:
|
|
continue # No evidence for this pair
|
|
|
|
weights = np.array([w_flip, w_conf, w_var, w_mp, w_cons])
|
|
if i_p==0 and j==0 and i_e==0:
|
|
logger.debug(f"Weights: {weights}, Evidences: {evidences}")
|
|
|
|
evidence_score = np.sum(weights * evidences)
|
|
|
|
context_evidence_scores.append(evidence_score)
|
|
|
|
if not context_evidence_scores:
|
|
# No high-confidence sources, use uniform weight
|
|
context_weight = 1.0
|
|
else:
|
|
# Average evidence across context
|
|
context_weight = np.mean(context_evidence_scores)
|
|
|
|
# Weight this prediction's LOGPROB by context reliability
|
|
weighted_logprobs.append(pred.raw_logprob_diff * context_weight)
|
|
total_weight += context_weight
|
|
|
|
# Aggregate: weighted mean in logprob space
|
|
if total_weight > 0:
|
|
final_logprob = sum(weighted_logprobs) / total_weight
|
|
else:
|
|
final_logprob = 0.0 # Fallback: neutral
|
|
|
|
# Convert to prediction
|
|
final_prob = lpr2prob(final_logprob)
|
|
pred_label = 1 if final_prob > 0.5 else 0
|
|
true_label = ex["vanilla_label"]
|
|
|
|
if pred_label == true_label:
|
|
correct += 1
|
|
total += 1
|
|
|
|
# Calibration error
|
|
calibration_errors.append(abs(final_prob - true_label))
|
|
|
|
accuracy = correct / total if total > 0 else 0.0
|
|
avg_calibration_error = np.mean(calibration_errors) if calibration_errors else 1.0
|
|
|
|
return {
|
|
"accuracy": accuracy,
|
|
"calibration_error": avg_calibration_error,
|
|
"correct": correct,
|
|
"total": total,
|
|
}
|
|
|
|
|
|
# Grid search (coarse)
|
|
best_acc = 0.0
|
|
best_calibration = 1.0
|
|
best_weights = None
|
|
best_metrics = None
|
|
|
|
# Try different weight combinations
|
|
# TODO try a proper grid
|
|
weight_grid = []
|
|
for w_flip in [0, 0.1, 0.5]:
|
|
for w_conf in [0, 0.1, 0.5]:
|
|
for w_var in [0, 0.1, 0.5]:
|
|
for w_mp in [0, 0.1, 0.5]:
|
|
for w_cons in [0, 0.1, 0.5]:
|
|
weights = np.array([w_flip, w_conf, w_var, w_mp, w_cons])
|
|
if weights.sum() > 0:
|
|
weights = weights / weights.sum().clip(min=1e-6) # Normalize to sum to 1
|
|
weight_grid.append(weights)
|
|
|
|
|
|
|
|
|
|
print("\n=== Weight Grid Search ===")
|
|
|
|
data = []
|
|
for weights in tqdm(weight_grid):
|
|
w_flip, w_conf, w_var, w_mp, w_cons = weights
|
|
metrics = evaluate_weights(
|
|
predictions, evidence_dict, test_group, w_flip, w_conf, w_var, w_mp, w_cons
|
|
)
|
|
|
|
|
|
if (metrics["accuracy"] >= best_acc) and (metrics["calibration_error"] < best_calibration):
|
|
best_acc = metrics["accuracy"]
|
|
best_calibration = metrics["calibration_error"]
|
|
best_weights = weights
|
|
best_metrics = metrics
|
|
|
|
data.append({
|
|
"weights": weights,
|
|
"accuracy": metrics["accuracy"],
|
|
"calibration_error": metrics["calibration_error"],
|
|
"correct": metrics["correct"],
|
|
"total": metrics["total"],
|
|
})
|
|
|
|
df_results = pd.DataFrame(data).sort_values(by=["accuracy", "calibration_error"], ascending=[False, True])
|
|
print("\n=== Sorted Grid Search Results ===")
|
|
print(df_results.head(30))
|
|
|
|
s_top_weights = pd.DataFrame(np.stack(df_results.head(30).weights.values), columns=['flip_contrib', 'conf', 'var', 'mp', 'cons']).mean().sort_values()
|
|
print("\n=== Top 30 Weights Average ===")
|
|
print(s_top_weights)
|
|
|
|
# TODO label best weights with df_evidence.columns
|
|
print(
|
|
f"\nBest weights: {best_weights} -> acc={best_acc:.3f}, cal_err={best_metrics['calibration_error']:.3f}"
|
|
)
|
|
|
|
# %% [code]
|
|
# Aggregate statistics with BEST weights
|
|
print("\n=== Ensemble Statistics (Best Weights) ===")
|
|
for uid, scores in target_scores.items():
|
|
mean = np.mean(scores)
|
|
std = np.std(scores)
|
|
# Find ground truth
|
|
true_label = next(ex["vanilla_label"] for ex in test_group if ex["uid"] == uid)
|
|
predicted_label = 1 if mean > 0.5 else 0
|
|
correct = "✓" if predicted_label == true_label else "✗"
|
|
uid_str = str(uid)[-6:] if isinstance(uid, (int, str)) else str(uid)[:6]
|
|
print(
|
|
f"{uid_str}: mean={mean:.3f}, std={std:.3f}, pred={predicted_label}, true={true_label} {correct}"
|
|
)
|
|
|
|
# %% [code]
|
|
# Save predictions as JSONL
|
|
output_path = Path("outputs/icm/evidence_test_predictions.jsonl")
|
|
output_path.parent.mkdir(parents=True, exist_ok=True)
|
|
|
|
with open(output_path, "w") as f:
|
|
for pred in predictions:
|
|
record = {
|
|
"target_uid": pred.target_uid,
|
|
"target_idx": pred.target_idx,
|
|
"score": pred.score,
|
|
"raw_logprob_diff": pred.raw_logprob_diff,
|
|
"context": [
|
|
{"uid": uid, "label": int(lbl), "raw_logprob": raw_lp, "flipped": flip}
|
|
for uid, lbl, raw_lp, flip in pred.context
|
|
],
|
|
"variations": pred.variations,
|
|
}
|
|
f.write(json.dumps(record) + "\n")
|
|
|
|
logger.info(f"Saved predictions to {output_path}")
|
|
|
|
# %% [markdown]
|
|
# ## Results Summary
|
|
#
|
|
# **Grid Search Findings (After Fixes):**
|
|
# - Best weights vary by run due to stochastic ensemble sampling
|
|
# - Top results often: pure flip_contrib [0.83, 0, 0.17, 0, 0] or pure variance [0, 0, 1, 0, 0]
|
|
# - Accuracy range: 63-91% depending on random flips/orderings (vs 50% random baseline)
|
|
# - Multiple weight combos often tie → suggests small budget (36 preds) creates sparse evidence
|
|
#
|
|
# **Key Insights:**
|
|
# 1. **Ensemble variance (epistemic uncertainty) is reliable signal** - low var = confident
|
|
# 2. **Flip_contrib (directional evidence) matters** - positive delta = flip improved coherence
|
|
# 3. **Proper consistency rules essential** - paraphrases agree, contradictions oppose
|
|
# 4. **Multi-flip processing densifies evidence graph** - not just first flip
|
|
#
|
|
# **Why results vary across runs?**
|
|
# - Budget=36 on 12 examples = ~3 predictions/example (sparse)
|
|
# - Random sampling of target, context, flips, orderings
|
|
# - No seed was set initially (now fixed with RANDOM_SEED=42)
|
|
# - Small group size (12) limits consistency constraints
|
|
#
|
|
# **Remaining TODOs:**
|
|
# 1. **Scale budget to 100+** predictions for denser evidence graph
|
|
# 2. **Add global examples** (2-3 from other groups) to break echo chambers
|
|
# 3. **Use Pearson correlation** for mutual_predictability (not just mean proximity)
|
|
# 4. **Test on full ICM** integration via simple_icm.py with larger dataset
|
|
|
|
# %% [markdown]
|
|
# ## Analysis & Next Steps
|
|
#
|
|
# **Observations:**
|
|
# - Evidence weights capture flip sensitivity—high weights = strong coupling between examples
|
|
# - Ensemble variance reveals aleatoric uncertainty (some examples inherently ambiguous)
|
|
# - Calibration: mean scores closer to 0/1 than raw logprobs (sigmoid normalization helps)
|
|
#
|
|
# **Refinements Needed:**
|
|
# 1. **Consistency weighting**: Currently hardcoded to 1, should check group rules (paraphrases agree, contradictions oppose)
|
|
# 2. **Global examples**: Mix 2 random examples per group to prevent local echo chambers
|
|
# 3. **Adaptive budget**: Spend more on high-variance targets
|
|
# 4. **Directional evidence**: Track if flip improves or worsens global energy
|
|
#
|
|
# **Theory Clarification (Epistemic vs Aleatoric):**
|
|
# - **Epistemic**: Reducible via more/better context → measured by variance across ensemble (low var = model "knows")
|
|
# - **Aleatoric**: Irreducible ambiguity in data → surfaces as consistency violations (e.g., paraphrases disagree)
|
|
# - **Our method**: Ensemble variance ≈ epistemic, consistency failures ≈ aleatoric, evidence weights ≈ structural/relational confidence
|
|
#
|
|
# **Concise Display Idea:**
|
|
# The `Prediction.__repr__` shows `P(uid | ctx1=lbl1*, ctx2=lbl2) = 0.75` where `*` marks flips.
|
|
# Could extend to a class that auto-formats for logging/notebooks.
|