mirror of
https://github.com/wassname/Unsupervised-Elicitation.git
synced 2026-10-30 13:00:17 +08:00
561 lines
19 KiB
Python
561 lines
19 KiB
Python
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)
|