wassname/Unsupervised-Elicitation
1
0
Fork 0
mirror of https://github.com/wassname/Unsupervised-Elicitation.git synced 2026-08-20 12:10:10 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix inconsistents

Browse Source 
fix inconsistents as combinatorial with no llm calls
This commit is contained in:
wassname
2025-10-05 09:10:22 +08:00
parent 65fac40c12
commit e4baa1b6cd
6 changed files with 157 additions and 677 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+1
View File
@@ -1,5 +1,6 @@
.env
.anycache/
docs/
# Byte-compiled / optimized / DLL files
__pycache__/
README.md
+16 -1
View File
@@ -1,9 +1,24 @@
This is a simplified fork of unsupervised-elicitation which is a paper that uses in-context metalearning, requiring it to be consistent and mutually predictive. We are using two things here
Implicit:
- in-context metalearning e.g. "given these examples complete the pattern"
Extinsic:
- consistancy
- mutual predictability
- simualted an
- nealing
Major changes:
- No leading the witness by using "find the truth" or "which is helpful" in the prompt
- ~~No "cheating" with consistency groups (e.g. one of these must be true)~~ hard to make it work without it
Fork to
- [x] refactor to UV
- [x] use openrouter logprob moddels
- [x] and simplify
- [x] replicate ![](docs/icm_progress.png)
- [ ] remove bias from the prompts "find truth" "which is helpfull" is too leading for unsupervised IMO
- [x] remove bias from the prompts "find truth" "which is helpfull" is too leading for unsupervised IMO
- [ ] add moral datasets e.g. daily dilemmas, ETHICS, Machiavelli, moral foundations vignettes
docs/original_ICM.py
-560
View File
@@ -1,560 +0,0 @@
import asyncio
import json
import math
import os
import random
from collections import Counter
from copy import deepcopy
from tqdm import tqdm
import numpy as np
from datasets import load_dataset
import argparse
from core.llm_api.llm import ModelAPI
from core.utils import setup_environment
from src.experiments.ICM_tools import (
propose_consistencyfix,
run_consistencyfix,
pick_two_inconsistent_claims,
update_assign_based_on_decision,
)
from src.model_querying.prompt_creation import (
get_decision_prompt,
get_judge_prompt_fewshot,
)
from src.model_querying.solution_extraction import (
extract_claim_logprobs,
extract_decision_logprobs,
)
from src.pipeline.pipeline import Pipeline, PipelineConfig
from src.tools.dataloaders import (
load_assignments,
load_problems_from_json,
load_problems_from_json_ids,
)
from src.tools.path_utils import get_default_results_directory, get_root_directory
def calculate_accuracy(train_data, inconsistent_pairs):
train_probs = []
for i in train_data.values():
if i["label"] is None:
continue
if i["label"] == 1:
train_probs.append(i["score"])
else:
train_probs.append(-i["score"])
if len(train_probs) == 0:
train_prob = 0
else:
train_prob = np.mean(train_probs)
return {
"train_accuracy": 0
if len(train_data) == 0
else np.mean([i["label"] == i["vanilla_label"] for i in train_data.values()]),
"train_label_distribution": Counter(
[i["vanilla_label"] for i in train_data.values()]
),
"train_predict_distribution": Counter(
[i["label"] for i in train_data.values()]
),
"train_prob": train_prob,
"train_size": len(train_data),
"inconsistent_num": len(inconsistent_pairs),
}
def update_assign(data):
for key, value in data.items():
if value["score"] > 0:
value["label"] = 1
else:
value["label"] = 0
return data
def fix_inconsistency(demonstrations, cur_metric, name, alpha, iter=0, K=20):
backup_metric = deepcopy(cur_metric)
if cur_metric["inconsistent_num"] == 0:
return demonstrations, cur_metric
cur_pool = {k: v for k, v in demonstrations.items() if v["label"] is not None}
assignment = cur_pool
best_metric = cur_metric
best_assignment = assignment
best_decision_id = None
for k in range(K):
pipeline = propose_consistencyfix(
args.model,
name=name,
iter=f"{iter}-{k}",
assignment=assignment,
)
results = asyncio.run(pipeline.run())
decisions = results["decisions"]
assignment = results["get_assign"]
for decision_id, decision in enumerate(decisions.values()):
tmp_decision_metric_list = []
tmp_decision_assignment_list = []
for score_idx, score in enumerate([0, 1]):
tmp_decision = deepcopy(decision)
tmp_decision["score"] = score
tmp_assignment = update_assign_based_on_decision(
deepcopy(assignment), tmp_decision
)
tmp_pipeline = run_consistencyfix(
model=args.model,
name=name,
iter=f"{iter}-{k}-{decision_id}-{score_idx}",
assignment=tmp_assignment,
)
tmp_results = asyncio.run(tmp_pipeline.run())
tmp_metric = tmp_results["evaluate"]
tmp_decision_metric_list.append(tmp_metric)
tmp_decision_assignment_list.append(tmp_assignment)
tmp_best_decision_id = np.argmax(
[get_energy(i, args.alpha) for i in tmp_decision_metric_list]
)
tmp_assignment = tmp_decision_assignment_list[tmp_best_decision_id]
tmp_metric = tmp_decision_metric_list[tmp_best_decision_id]
if get_energy(tmp_metric, args.alpha) >= get_energy(best_metric, args.alpha):
best_decision_id = decision_id
best_metric = tmp_metric
best_assignment = tmp_assignment
break
if best_decision_id is None:
break
elif best_metric["inconsistent_num"] == 0:
assignment = best_assignment
break
else:
assignment = best_assignment
for k in assignment:
demonstrations[k] = assignment[k]
return demonstrations, best_metric
def get_pipeline(
model,
name=None,
use_cache=True,
num_problems=None,
decision_id=None,
iter=None,
assignment=None,
):
pipeline_name = f"iterative-truth-assign-iter-{iter}"
if decision_id is not None:
pipeline_name += f"-{decision_id}"
if name is not None:
pipeline_name += "-" + name
ROOT_DIR = get_root_directory()
DATA_DIR = ROOT_DIR / "data"
pipeline_config = PipelineConfig(
pipeline_name,
anthropic_num_threads=40,
openai_fraction_rate_limit=0.99,
num_problems=num_problems,
use_cache=use_cache,
)
pipeline = Pipeline(pipeline_config)
assert assignment is not None
initial_assign = pipeline.add_load_data_step(
"get_assign", load_assignments, assignment
)
def add_train_demonstrations(train_data):
copy_data = deepcopy(train_data)
copy_data = {k: v for k, v in copy_data.items() if v["label"] is not None}
keys = list(copy_data.keys())
values = list(copy_data.values())
saved_keys = [
"prompt",
"question",
"choice",
"choice_2",
"consistency_id",
"consistency_key",
"source",
"label",
"vanilla_label",
]
values = []
for i in copy_data.values():
values.append({saved_key: i[saved_key] for saved_key in saved_keys if saved_key in i})
for idx, key in enumerate(keys):
tmp_keys, tmp_values = [], []
for j, (prev_key, prev_value) in enumerate(zip(keys, values)):
if j != idx:
tmp_keys.append(prev_key)
tmp_values.append(prev_value)
demos = {
prev_key: prev_value
for j, (prev_key, prev_value) in enumerate(zip(tmp_keys, tmp_values))
}
sorted_demos = {}
for k, v in demos.items():
q = v["consistency_id"]
if q not in sorted_demos:
sorted_demos[q] = []
sorted_demos[q].append((k, v))
out_sorted_demos = {}
for group in sorted_demos.values():
for k, v in group:
out_sorted_demos[k] = v
copy_data[key]["demonstration"] = out_sorted_demos
return copy_data
merged_train_data = pipeline.add_transformation_step(
"add_train_demonstration",
add_train_demonstrations,
dependencies=[initial_assign],
)
get_train_preds = pipeline.add_query_step(
"get_train_preds",
model,
get_judge_prompt_fewshot,
extract_claim_logprobs,
dependencies=[merged_train_data],
logprobs=20,
max_tokens=1,
use_cache=use_cache,
)
pick_claims = pipeline.add_transformation_step(
"pick_two_inconsistent_claims",
pick_two_inconsistent_claims,
dependencies=[initial_assign],
)
eval_preds = pipeline.add_eval_step(
"evaluate",
calculate_accuracy,
dependencies=[get_train_preds, pick_claims],
)
return pipeline
async def predict_assignment(model, example, demonstrations):
demos = [
v
for k, v in demonstrations.items()
if k != example["uid"] and v["label"] is not None
]
anthropic_requests = [
model_api(
model,
get_judge_prompt_fewshot(
example,
demos,
pipeline=False,
),
logprobs=20,
max_tokens=1,
parse_fn=extract_claim_logprobs,
)
]
responses = await asyncio.gather(*anthropic_requests)
score = responses[0][0]["score"]
new_label = score > 0
return int(new_label)
def get_temperature(
iteration, initial_temp, final_temp, decay_rate, schedule="exp"
):
"""
Calculate the temperature for simulated annealing.
Parameters:
- iteration: Current iteration number.
- initial_temp: Initial temperature.
- decay_rate: Rate at which the temperature decreases.
Returns:
- Current temperature.
"""
if schedule == "exp":
return max(final_temp, initial_temp * (decay_rate**iteration))
elif schedule == "log":
return max(final_temp, initial_temp / (1 + 2 * np.log(1 + iteration)))
else:
assert False
def get_energy(metric, alpha):
return alpha * metric["train_prob"] - metric["inconsistent_num"]
def get_args():
parser = argparse.ArgumentParser()
parser.add_argument("--alpha", type=float, default=30)
parser.add_argument("--seed", type=int, default=27565976)
parser.add_argument("--testbed", type=str, default="gsm8k")
parser.add_argument("--model", type=str, default="meta-llama/Llama-3.1-70B")
parser.add_argument("--batch_size", type=int, default=256)
parser.add_argument("--num_seed", type=int, default=8)
parser.add_argument("--K", type=int, default=3000)
parser.add_argument("--consistency_fix_K", type=int, default=10)
parser.add_argument("--decay", type=float, default=0.99)
parser.add_argument("--initial_T", type=float, default=10)
parser.add_argument("--final_T", type=float, default=0.01)
parser.add_argument("--scheduler", type=str, default="log")
args = parser.parse_args()
return args
def load_data(args):
if args.testbed == "alpaca":
with open(get_root_directory() / "data/train_alpaca.json") as f:
train = json.load(f)
template = """Human: {question}
Response 1: {choice}
Response 2: {choice_2}
Claim: Response 1 is more helpful and harmless than Response 2.
I think this claim is """
for i in train:
i['source'] = 'alpaca'
i['consistency_key'] = 'A>B' if i['label'] else 'B>A'
i['prompt'] = template.format(question=i['question'], choice=i['choice'].replace("\n", ""), choice_2=i['choice_2'].replace("\n", ""))
args.GROUP_SIZE = 2
elif args.testbed == 'gsm8k':
with open(get_root_directory() / "data/train_gsm8k.json") as f:
train = json.load(f)
template = """Question: {question}
Claim: {answer}
I think this claim is """
for i in train:
i['source'] = 'gsm8k'
i['consistency_key'] = i['answer']
i['prompt'] = template.format(question=i['question'], answer=i['choice'])
args.GROUP_SIZE = 4
elif args.testbed == 'truthfulQA':
with open(get_root_directory() / "data/train_truthfulqa.json") as f:
train = json.load(f)
template = """Question: {question}
Claim: {answer}
I think this claim is """
for i in train:
i['source'] = 'truthfulQA'
i['consistency_key'] = 'A' if i['label'] else 'B'
i['prompt'] = template.format(question=i['question'], answer=i['choice'])
args.GROUP_SIZE = 4
elif args.testbed == 'truthfulQA-preference':
with open(get_root_directory() / "data/train_truthfulqa_preference.json") as f:
train = json.load(f)
template = """Question: {question}
Answer 1: {choice}
Answer 2: {choice_2}
Claim: Answer 1 is more truthful than Answer 2.
I think this claim is """
for i in train:
i['source'] = 'truthfulQA-preference'
i['consistency_key'] = 'A>B' if i['label'] else 'B>A'
i['prompt'] = template.format(question=i['question'], choice=i['choice'], choice_2=i['choice_2'])
args.GROUP_SIZE = 2
train_map = {}
for i in train:
if i['consistency_id'] not in train_map:
train_map[i['consistency_id']] = []
train_map[i['consistency_id']].append(i)
out = []
for key in train_map:
out += train_map[key]
train = out
# sample a batch of batch_size datapoints
fewshot_ids = random.sample(
list(range(len(train)// args.GROUP_SIZE)), args.batch_size // args.GROUP_SIZE
)
fewshot_ids = [
i * args.GROUP_SIZE + j for i in fewshot_ids for j in range(args.GROUP_SIZE)
]
return train, fewshot_ids
def initialize(train, fewshot_ids, args):
demonstrations = {}
unlabeled_ids = []
whole_ids = []
seed_ids = []
random_init_labels = [1] * (args.num_seed // 2) + [0] * (args.num_seed // 2)
random.shuffle(random_init_labels)
for id, i in enumerate(fewshot_ids):
item = train[i]
item["vanilla_label"] = item["label"] # store dataset labels to measure agreement during the searching process
item["uid"] = id
whole_ids.append(item["uid"])
if id >= args.num_seed: # set labels to None
item["label"] = None
item["type"] = "predict"
unlabeled_ids.append(item["uid"])
else: # set random labels
item["type"] = "seed"
item["label"] = random_init_labels[id]
seed_ids.append(item["uid"])
demonstrations[id] = item
return demonstrations, unlabeled_ids, whole_ids, seed_ids
def main(args):
train, fewshot_ids = load_data(args)
demonstrations, unlabeled_ids, whole_ids, seed_ids = initialize(train, fewshot_ids, args)
cur_metric = {
"train_prob": -1e6,
"inconsistent_num": 100000,
"train_accuracy": 1.0,
"train_predict_distribution": {"0": 0, "1": 0},
"train_label_distribution": {"0": 0, "1": 0},
}
print('init random labels = ', Counter([i['label'] for i in demonstrations.values() if i['type'] == 'seed']), 'init label acc = ', np.mean([i['label'] == i['vanilla_label'] for i in demonstrations.values() if i['type'] == 'seed']))
name = f"{args.testbed}-llama70b-K{args.K}-bc{args.batch_size}_seed{args.seed}-initialsize{args.num_seed}-weighted{args.alpha}-decay{args.decay}-initialT{args.initial_T}-finalT{args.final_T}-scheduler{args.scheduler}"
iter = 0
flip_cnt = 0
example_id = 0
for _ in tqdm(range(args.K), desc="searching"):
cur_pool = {
k: v for k, v in demonstrations.items() if v["label"] is not None
}
initial_demos = deepcopy(demonstrations)
if iter == 0:
pipeline = get_pipeline(
args.model,
name=name,
num_problems=None,
iter=iter,
assignment=cur_pool,
)
results = asyncio.run(pipeline.run())
cur_metric = results["evaluate"]
demonstrations, cur_metric = fix_inconsistency(
demonstrations, cur_metric, name, args.alpha, iter=iter, K=args.consistency_fix_K
)
cur_pool = {
k: v for k, v in demonstrations.items() if v["label"] is not None
}
while True: # weighted sampling
candidates_ids = whole_ids
weights = [1 for _ in range(len(candidates_ids))]
for i in candidates_ids:
if i in cur_pool:
same_consistency_group_ids = [j for j in candidates_ids if demonstrations[j]["consistency_id"] == demonstrations[i]["consistency_id"]]
for j in same_consistency_group_ids:
if j not in cur_pool:
weights[j] = 100
example_id = random.choices(candidates_ids, k=1, weights=weights)[0]
break
new_label = asyncio.run(
predict_assignment(
args.model,
demonstrations[example_id],
cur_pool,
)
)
if demonstrations[example_id]["label"] != new_label:
tmp_demonstrations = deepcopy(demonstrations)
tmp_demonstrations[example_id]["label"] = new_label
dummy_metric = {
"train_prob": -1e6,
"inconsistent_num": 100000,
"train_accuracy": 1.0,
"train_predict_distribution": {"0": 0, "1": 0},
"train_label_distribution": {"0": 0, "1": 0},
}
tmp_demonstrations, _ = fix_inconsistency(
tmp_demonstrations,
dummy_metric,
name + "newlabelexplore",
args.alpha,
iter=iter,
K=10,
)
tmp_pool = {
k: v
for k, v in tmp_demonstrations.items()
if v["label"] is not None
}
pipeline = get_pipeline(
model=args.model,
name=name,
num_problems=None,
iter=iter,
assignment=tmp_pool,
)
results = asyncio.run(pipeline.run())
metric = results["evaluate"]
T = get_temperature(
flip_cnt, args.initial_T, args.final_T, args.decay, schedule=args.scheduler
)
print(f"iter = {iter}, pool size = {len(cur_pool)}, cur acc = {cur_metric['train_accuracy']}, new acc = {metric['train_accuracy']}, cur score = {get_energy(cur_metric, args.alpha)}, new score = {get_energy(metric, args.alpha)}, cur inconsistent num = {cur_metric['inconsistent_num']}, new inconsistent num = {metric['inconsistent_num']}")
print('cur label distribution = ', Counter([i['label'] for i in demonstrations.values() if i['label'] is not None]))
print('new label distribution = ', Counter([i['label'] for i in tmp_demonstrations.values() if i['label'] is not None]))
accept_prob = math.exp((get_energy(metric, args.alpha) - get_energy(cur_metric, args.alpha)) / T)
print("accept prob = ", accept_prob)
if random.random() < accept_prob:
print("accept")
demonstrations = tmp_demonstrations
flip_cnt += 1
cur_metric = metric
with open(f"log_{name}.jsonl", "a") as f:
f.write(json.dumps({
"iter": iter,
"flip_cnt": flip_cnt,
"acc": cur_metric['train_accuracy'],
"score": get_energy(cur_metric, args.alpha),
}) + "\n")
else:
print("reject")
print("=" * 100)
iter += 1
if __name__ == "__main__":
setup_environment(logger_level="error")
model_api = ModelAPI(anthropic_num_threads=20, openai_fraction_rate_limit=0.99)
args = get_args()
print("task: ", args.testbed)
random.seed(args.seed)
main(args)
pyproject.toml
+5
View File
@@ -9,6 +9,7 @@ authors = [
requires-python = ">=3.10"
dependencies = [
"adjusttext>=1.3.0",
"aiocache>=0.12.3",
"alembic>=1.16.5",
"altair>=5.5.0",
"anthropic>=0.69.0",
@@ -51,6 +52,10 @@ unsupervised-elicitation = "unsupervised_elicitation:main"
requires = ["setuptools>=61"]
build-backend = "setuptools.build_meta"
[tool.setuptools.packages.find]
where = ["."] # search the root directory
include = ["src*"]
[tool.uv.sources]
openrouter-wrapper = { path = "../llm_moral_lb_v2/openrouter_wrapper", editable = true }
nbs/simple_icm.py → src/simple_icm.py
+124 -116
View File
@@ -13,16 +13,17 @@ import matplotlib.pyplot as plt
from pathlib import Path
import math
import numpy as np
from anycache import anycache
from functools import lru_cache
import os, sys
import pandas as pd
from dataclasses import dataclass, asdict
import dotenv
from loguru import logger
from openrouter_wrapper.logprobs import openrouter_completion_wlogprobs, get_logprobs_choices, LogprobsNotSupportedError # User's wrapper
from typing import List, Tuple
from typing import List, Tuple, Callable
import asyncio
from aiocache import cached
from itertools import combinations
from copy import deepcopy
try:
from IPython import get_ipython
@@ -40,6 +41,9 @@ dotenv.load_dotenv()
logger.remove()
logger.add(sys.stderr, format="<green>{time:YYYY-MM-DD HH:mm}</green> | <level>{level}</level> | <cyan>{message}</cyan>", colorize=True, level="DEBUG")
# Global cost tracker
total_cost = 0.0
# %% [code]
@dataclass
@@ -50,8 +54,8 @@ class Config:
decay_rate: float = 0.99
beta: float = 2.0
num_seed: int = 8
max_iters: int = 950 # Small for demo; increase for more
n_shots: int = 8 # Number of in-context examples
max_iters: int = 2500000 # should be at least dataset size X 2
n_shots: int = 6 # Number of in-context examples
model_id: str = "meta-llama/llama-3.1-8b-instruct" # Logprobs supported
provider_whitelist: Tuple[str] = None # None to let OpenRouter choose
out_dir: Path = Path("../outputs/icm") # Directory to save outputs
@@ -74,40 +78,21 @@ C.out_dir.mkdir(parents=True, exist_ok=True)
# )
logger.info(f"Config: {C}")
config_dict = asdict(C)
config_dict['out_dir'] = str(config_dict['out_dir'])
with open(C.out_dir / "icm_config.json", "w") as f:
json.dump(config_dict, f, indent=2)
# %% [code]
from datasets import load_dataset
from src.data.truthfulqa import load_truthfulqa, is_consistent
# Load larger HF dataset: Yik/truthfulQA-bool
dataset = load_dataset("wassname/truthful_qa_v2", split="validation", name="binary")
# semantic_anchor = "0 is false, 1 is true."
data = []
group_id = 0
groups = {}
for idx, item in enumerate(dataset):
claim = f"Question: {item['Question']}\nClaim: {item['Answer']}"
label = item['label']
consistency_id = item['question_id']
example = {
'uid': idx,
'prompt': claim,
'vanilla_label': label,
'consistency_id': consistency_id,
'consistency_key': 'A' if label == 1 else 'B'
}
data.append(example)
# HACK: Limit to small batch for demo
# data = data[:512]
logger.info(f"Loaded {len(data)} examples from TruthfulQA-bool")
data = load_truthfulqa()
logger.info("Loaded {} examples", len(data))
# %% [code]
# Initialize: Random labels for first num_seed, None for others
def initialize_data(data, config):
demonstrations = {item['uid']: item.copy() for item in data}
demonstrations = {item['uid']: deepcopy(item) for item in data}
labeled_uids = random.sample(list(demonstrations.keys()), min(config.num_seed, len(data)))
for uid in demonstrations:
demonstrations[uid]['label'] = None
@@ -125,35 +110,40 @@ logger.info("Initialized labels: {}", {k: v['label'] for k, v in demonstrations.
def print_messages(messages):
return "\n".join([f"**{m['role'].upper()}**: {m['content']}" for m in messages])
@cached()
async def cached_openrouter_completion_wlogprobs(*args, **kwargs):
# TODO: Async caching (e.g., with aiocache)
return await openrouter_completion_wlogprobs(*args, **kwargs)
async def predict_label(example_uid, current_demos, config=C, verbose=False):
async def predict_label(example_uid, current_demos, config=C, verbose=False, all_demos=None):
# simplified from https://github.com/Jiaxin-Wen/Unsupervised-Elicitation/blob/master/src/experiments/ICM.py#L254
# Sort by consistency_key for relevance, limit to 16
relevant_demos = sorted(
[demo for uid, demo in current_demos.items() if uid != example_uid and demo['label'] is not None],
key=lambda x: x.get('consistency_key', '')
)[:C.n_shots]
# TODO: Use all/subsampled demos like original (group/sort by consistency_id for relevance), handling context limits (high value, med complexity).
global total_cost
# Group and interleave demos by consistency_id like original
grouped_demos = {}
for uid, demo in current_demos.items():
if uid != example_uid and demo['label'] is not None:
grouped_demos.setdefault(demo['consistency_id'], []).append(demo)
# FIXME: some LLM's have a positional bias, so we should randomize order
relevant_demos = []
for group in grouped_demos.values():
random.shuffle(group) # Reduce order bias
relevant_demos.extend(group[:2]) # Limit per group for balance
instruction = "Carefully examine the noisy example pairs of (Candidate, Set). Then as consistently and coherently as possible complete the pattern, you may provide reasoning after.\n\n"
# if verbose>1:
# instruction = "Examine the examples below and determine the category 0 or 1 for the final claim based on the pattern. After 0 or 1 output a single word to describe how category 1 compares to category 0.\n\n"
relevant_demos = relevant_demos[:C.n_shots] # Cap total
instruction = "Examine these examples and complete the pattern for the new one.\n\n"
fewshot = []
# FIXME to not lead the unsupervised model, we should avoid true/false or even 0/1 and try to use neutral labels like A/B or similar
for idx, demo in enumerate(relevant_demos):
for demo in relevant_demos:
label_str = "A" if demo['label'] == 1 else "B"
fewshot.append(f"\nCandidate: {demo['prompt']}\nSet: {label_str}\n")
fewshot.append(f"\n\n## Candidate:\n{demo['prompt']}\n## Set: {label_str}")
target_prompt = demonstrations[example_uid]['prompt']
# Use all_demos if provided (for unlabeled examples), otherwise use current_demos
demos_source = all_demos if all_demos is not None else current_demos
target_prompt = demos_source[example_uid]['prompt']
messages = [
{"role": "user", "content": instruction+"".join(fewshot)+f"Candidate: {target_prompt}\n"},
{"role": "assistant", "content": "\n\nSet:"} # Assistant prefill to ensure
{"role": "user", "content": instruction+"".join(fewshot)+f"\n\n## Candidate:\n{target_prompt}"},
{"role": "assistant", "content": "\n## Set:"} # Assistant prefill to ensure
]
@@ -161,11 +151,13 @@ async def predict_label(example_uid, current_demos, config=C, verbose=False):
model_id=config.model_id,
provider_whitelist=config.provider_whitelist,
messages=messages,
max_completion_tokens=90 if verbose else 5,
max_completion_tokens=160 if verbose else 5,
temperature=0.4,
top_logprobs=8,
)
total_cost += response.get('usage', {}).get('cost', 0.0)
if verbose:
logger.info(f"Debug Prediction - UID {example_uid}:")
logger.info(f"messages: {print_messages(messages)}")
@@ -229,7 +221,7 @@ def compute_energy(demos, config=C):
if len(set(all_labels)) < len(key_groups): # Not all oppose if num unique labels < num keys
num_inconsistent += max(0, len(items) - len(set(all_labels)))
energy = config.alpha * avg_lprob - num_inconsistent
energy = config.alpha * avg_lprob - num_inconsistent - (num_inconsistent / max(1, len(labeled))) # Normalized penalty
accuracy = np.mean([d['label'] == d['vanilla_label'] for d in labeled])
return energy, {
'avg_lprob': avg_lprob,
@@ -241,60 +233,77 @@ def compute_energy(demos, config=C):
logger.info("Initial energy: {}", compute_energy(demonstrations))
# %% [code]
async def fix_inconsistencies_simple(demos, config=C, max_fixes=5):
"""Simple consistency fix: for inconsistent pairs, enumerate label combos, re-predict, pick max energy."""
def get_kflip_neighbors(group_uids, demos, k):
"""
Generate all label assignments that are k flips away from current.
Returns list of [(uid, new_label), ...] tuples.
"""
labeled_uids = [uid for uid in group_uids if demos[uid]['label'] is not None]
neighbors = []
for combo in combinations(labeled_uids, k):
flips = [(uid, 1 - demos[uid]['label']) for uid in combo]
neighbors.append(flips)
return neighbors
async def fix_inconsistencies_greedy(demos, config=C, max_fixes=20, max_flips=3, is_consistent: Callable = is_consistent):
"""
Greedy consistency fix: Try k-flip neighborhoods (k=1,2,...) until we find
a consistent assignment that improves energy.
"""
for fix_iter in range(max_fixes):
# Find inconsistent pairs
# Find inconsistent groups
groups = {}
for uid, demo in demos.items():
if demo['label'] is not None:
cid = demo['consistency_id']
if cid not in groups:
groups[cid] = []
groups[cid].append(uid)
groups.setdefault(demo['consistency_id'], []).append(uid)
# Find first inconsistent pair
inconsistent_pair = None
# Find first inconsistent group
inconsistent_group = None
for cid, uids in groups.items():
labels = [demos[uid]['label'] for uid in uids if demos[uid]['label'] is not None]
if len(set(labels)) > 1: # Inconsistent
inconsistent_pair = (uids[0], uids[1])
if not is_consistent(uids, demos):
inconsistent_group = uids
break
if inconsistent_pair is None:
break # No more inconsistencies or not yet enough labels to have any effect
if not inconsistent_group:
break # All consistent!
uid1, uid2 = inconsistent_pair
old_energy, _ = compute_energy(demos, config)
best_energy = old_energy
best_flips = None
# Enumerate all 4 label combinations and pick max energy
# FIXME doesn't this call predict_label twice per option?
options = [(0, 0), (0, 1), (1, 0), (1, 1)]
best_energy = float('-inf')
best_option = None
for label1, label2 in options:
temp_demos = demos.copy()
temp_demos[uid1]['label'] = label1
temp_demos[uid2]['label'] = label2
# Try k=1, 2, 3, ... flips until we find improvement
for k in range(1, min(max_flips + 1, len(inconsistent_group) + 1)):
neighbors = get_kflip_neighbors(inconsistent_group, demos, k)
# Re-predict labels with new context to update scores
current_labeled = {k: v for k, v in temp_demos.items() if v['label'] is not None}
new_label1, score1 = await predict_label(uid1, current_labeled, config)
new_label2, score2 = await predict_label(uid2, current_labeled, config)
temp_demos[uid1]['score'] = score1
temp_demos[uid2]['score'] = score2
for flips in neighbors:
# Apply flips temporarily
temp_demos = deepcopy(demos)
for uid, new_label in flips:
temp_demos[uid]['label'] = new_label
# Check if this is consistent
if not is_consistent(inconsistent_group, temp_demos):
continue # Skip inconsistent neighbors
# Score it
energy, _ = compute_energy(temp_demos, config)
if energy > best_energy:
best_energy = energy
best_flips = flips
energy, _ = compute_energy(temp_demos, config)
if energy > best_energy:
best_energy = energy
best_option = (label1, label2, score1, score2)
if best_flips:
break # Found improvement with k flips, don't try k+1
# Apply best option
demos[uid1]['label'] = best_option[0]
demos[uid1]['score'] = best_option[2]
demos[uid2]['label'] = best_option[1]
demos[uid2]['score'] = best_option[3]
# Apply best flips if improvement found
if best_flips and best_energy > old_energy:
for uid, new_label in best_flips:
demos[uid]['label'] = new_label
else:
break # No improvement possible, stop trying
return demos
@@ -305,7 +314,7 @@ async def run_icm(demonstrations, config=C):
accuracies = []
# Fix any initial inconsistencies from random initialization
demonstrations = await fix_inconsistencies_simple(demonstrations, config)
demonstrations = await fix_inconsistencies_greedy(demonstrations, config)
current_labeled = {k: v for k, v in demonstrations.items() if v['label'] is not None}
old_energy, old_metrics = compute_energy(demonstrations, config)
@@ -360,17 +369,21 @@ async def run_icm(demonstrations, config=C):
verbose = 2
else:
verbose = 0
new_label, score = await predict_label(example_uid, current_labeled, config, verbose=verbose)
new_label, score = await predict_label(example_uid, current_labeled, config, verbose=verbose, all_demos=demonstrations)
# Update with new label and fix any inconsistencies
temp_demos = demonstrations.copy()
# Update with new label and fix any inconsistencies ONLY if label changed
temp_demos = deepcopy(demonstrations)
temp_demos[example_uid]['label'] = new_label
temp_demos[example_uid]['score'] = score
temp_demos = await fix_inconsistencies_simple(temp_demos, config)
# Only fix inconsistencies if the label actually changed
if demonstrations[example_uid]['label'] != new_label:
temp_demos = await fix_inconsistencies_greedy(temp_demos, config)
# Compute new energy
new_energy, new_metrics = compute_energy(temp_demos, config)
delta = new_energy - old_energy
temp_demos[example_uid]['deltaE'] = delta
# Annealing decision
T = max(config.final_t, config.initial_t / (1 + config.beta * math.log(1 + iter)))
@@ -387,7 +400,7 @@ async def run_icm(demonstrations, config=C):
accuracies.append(new_metrics['accuracy'])
if iter % C.log_interval == 0:
logger.info(f"Progress: Labeled {new_metrics['num_labeled']}, Inconsistents: {new_metrics['num_inconsistent']}. Acc: {new_metrics['accuracy']:.2f}, Energy: {old_energy:.2f}")
logger.info(f"Progress: Labeled {new_metrics['num_labeled']}, Inconsistents: {new_metrics['num_inconsistent']}. Acc: {new_metrics['accuracy']:.2f}, Energy: {old_energy:.2f}. Cost so far: ${total_cost:.4f}")
return demonstrations, energies, accuracies
@@ -398,6 +411,7 @@ final_demos, energies, accuracies = asyncio.run(run_icm(demonstrations, C))
# Final metrics
final_energy, final_metrics = compute_energy(final_demos, C)
logger.info("\nFinal Results:")
logger.info("Total cost: ${:.4f}", total_cost)
logger.info("Energy: {:.2f}", final_energy)
# TODO show vanilla accuracy here for comparison
logger.info("Accuracy vs vanilla: {:.2f}, initial {:.2f}", final_metrics['accuracy'], accuracies[0])
@@ -408,19 +422,16 @@ logger.info("Inconsistencies: {}", final_metrics['num_inconsistent'])
df = pd.DataFrame(final_demos).T
df.to_parquet(C.out_dir / "icm_final_labels.parquet")
logger.info("\nFinal labels:")
for uid, demo in final_demos.items():
label = demo['label']
if label is not None:
logger.info(f"UID {uid} ({demo['consistency_id']}): {label} (vanilla: {demo['vanilla_label']})")
df_labeled = df.dropna(subset='label').sort_values(by='score', key=np.abs, ascending=False)
df_labeled_disagreed = df_labeled[df_labeled['vanilla_label'] != df_labeled['label']]
print(f"\nFinal labeled examples (total {len(df_labeled)}):")
print(df_labeled_disagreed[['consistency_id', 'label', 'vanilla_label', 'score', 'prompt']])
for uid, row in df_labeled_disagreed.iterrows():
print(f"\n## Candidate: {row['prompt']}\nICM Set: {'A' if row['label']==1 else 'B'}, Vanilla Set: {'A' if row['vanilla_label']==1 else 'B'}, score={row['score']}\n")
json.dump(
asdict(C),
open(C.out_dir / "icm_config.json", "w")
)
# TODO put them in file, only print a top few disagreements
# %% [code]
# Simple visualization (requires matplotlib)
@@ -443,27 +454,24 @@ plt.tight_layout()
plt.savefig("icm_progress.png")
plt.show()
# %% [markdown]
# ## Next Steps & Limitations
# Prioritized by complexity (low/med/high) vs. value (low/med/high) based on paper comparison. Limitations noted with potential fixes.
# # Existing TODOs:
# # Existing TODOs:
# - [x] Load larger HF dataset ('Yik/truthfulQA-bool' subset, formatted to messages).
# - [x] Refine few-shot prompt from original get_judge_prompt_fewshot.
# - [x] Add weighted sampling for inconsistent groups (no longer random).
# - [x] Remove biased instruction - pure pattern completion for unsupervised elicitation.
# - [x] Simplify consistency fix - enumerate label combos, re-predict, pick max energy (no LLM meta-reasoning).
# - [x] Add log temperature schedule: max(Tmin, T0 / (1 + β log(n))) - Low complexity, Med value: Better early exploration.
# - [ ] Add caching for predictions (dict/file-based) - Low complexity, Med value: @lru_cache already used, could add disk cache.
# - [ ] Implement async batch predictions in predict_label - Med complexity, Med value: Use asyncio.gather for concurrent API calls.
# - [ ] Full mutual predictability: Use all/subsampled demos in predict_label - High complexity, High value: Handle context limits.
# - [ ] Test robustness with worst-case init: Add golden/random/worst init options - Med complexity, Low value: Validate like paper Sec. 5.
# # - [ ] Dynamically select demos in few-shot (group/sort by consistency_id) - Med complexity, High value.
# # - [ ] Weight sampling by energy delta potential - Low complexity, Med value.
# #
# # Limitations (from Paper Sec. 9) & Potential Fixes:
# - [ ] Salient concepts only: ICM can't elicit non-salient private preferences (e.g., "sun" poems); fix: Combine with weak supervision.
# - [ ] Context length limits: Can't fit all N demos for large datasets; fix: Subsample relevant demos or use long-context models.
# - [ ] Degenerate solutions without consistency: Risk of all-same labels; mitigated by logical consistency term.
# - [ ] Inference cost: 2-3 fwd passes per point (paper App. B); fix: Caching + batching reduces API hits.
# - [ ] I'd like to record what it thinks the labels represent e.g. "misconception" "virtue" etc, and find which answer leads to the best energy.
# - [ ] consistency groups rely on known labels, this is a limitation for real unsupervised use
uv.lock
Generated
+11
View File
@@ -26,6 +26,15 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/53/1c/8feedd607cc14c5df9aef74fe3af9a99bf660743b842a9b5b1865326b4aa/adjustText-1.3.0-py3-none-any.whl", hash = "sha256:da23d7b24b6db5ffa039bb136bfa556207365e32f48ac74b07ad26dd485bc691", size = 13154, upload-time = "2024-10-31T16:45:35.227Z" },
]
[[package]]
name = "aiocache"
version = "0.12.3"
source = { registry = "https://pypi.org/simple" }
sdist = { url = "https://files.pythonhosted.org/packages/7a/64/b945b8025a9d1e6e2138845f4022165d3b337f55f50984fbc6a4c0a1e355/aiocache-0.12.3.tar.gz", hash = "sha256:f528b27bf4d436b497a1d0d1a8f59a542c153ab1e37c3621713cb376d44c4713", size = 132196, upload-time = "2024-09-25T13:20:23.823Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/37/d7/15d67e05b235d1ed8c3ce61688fe4d84130e72af1657acadfaac3479f4cf/aiocache-0.12.3-py2.py3-none-any.whl", hash = "sha256:889086fc24710f431937b87ad3720a289f7fc31c4fd8b68e9f918b9bacd8270d", size = 28199, upload-time = "2024-09-25T13:20:22.688Z" },
]
[[package]]
name = "aiohappyeyeballs"
version = "2.6.1"
@@ -3748,6 +3757,7 @@ version = "0.1.0"
source = { editable = "." }
dependencies = [
{ name = "adjusttext" },
{ name = "aiocache" },
{ name = "alembic" },
{ name = "altair" },
{ name = "anthropic" },
@@ -3793,6 +3803,7 @@ dev = [
[package.metadata]
requires-dist = [
{ name = "adjusttext", specifier = ">=1.3.0" },
{ name = "aiocache", specifier = ">=0.12.3" },
{ name = "alembic", specifier = ">=1.16.5" },
{ name = "altair", specifier = ">=5.5.0" },
{ name = "anthropic", specifier = ">=0.69.0" },