From e4baa1b6cdfc28b276d6cbf395c5a5153b4fb8e6 Mon Sep 17 00:00:00 2001
From: wassname <1103714+wassname@users.noreply.github.com>
Date: Sun, 5 Oct 2025 09:10:22 +0800
Subject: [PATCH] fix inconsistents
fix inconsistents as combinatorial with no llm calls
---
.gitignore | 1 +
README.md | 17 +-
docs/original_ICM.py | 560 -------------------------------------
pyproject.toml | 5 +
{nbs => src}/simple_icm.py | 240 ++++++++--------
uv.lock | 11 +
6 files changed, 157 insertions(+), 677 deletions(-)
delete mode 100644 docs/original_ICM.py
rename {nbs => src}/simple_icm.py (68%)
diff --git a/.gitignore b/.gitignore
index 2f2a760..ff21c0f 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,5 +1,6 @@
.env
.anycache/
+docs/
# Byte-compiled / optimized / DLL files
__pycache__/
diff --git a/README.md b/README.md
index 1afef52..c77c1ff 100644
--- a/README.md
+++ b/README.md
@@ -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 
-- [ ] 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
diff --git a/docs/original_ICM.py b/docs/original_ICM.py
deleted file mode 100644
index 7d8335f..0000000
--- a/docs/original_ICM.py
+++ /dev/null
@@ -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)
diff --git a/pyproject.toml b/pyproject.toml
index 81cba5e..8a463c3 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -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 }
diff --git a/nbs/simple_icm.py b/src/simple_icm.py
similarity index 68%
rename from nbs/simple_icm.py
rename to src/simple_icm.py
index a866eca..44e8bd3 100644
--- a/nbs/simple_icm.py
+++ b/src/simple_icm.py
@@ -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="{time:YYYY-MM-DD HH:mm} | {level} | {message}", 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
diff --git a/uv.lock b/uv.lock
index d9a0948..1bae939 100644
--- a/uv.lock
+++ b/uv.lock
@@ -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" },