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

working

Browse Source
This commit is contained in:
wassname
2025-10-02 22:14:40 +08:00
parent 6943102504
commit 1002829c7b
26 changed files with 502 additions and 3576 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+4 -4
View File
@@ -1,14 +1,14 @@
Fork to
- [ ] add moral datasets e.g. daily dilemmas, ETHICS, Machiavelli, moral foundations vignettes
- [x] refactor to UV
- [ ] and simplify
- [ ] replicate
- [x] and simplify
- [x] replicate ![](docs/icm_progress.png)
- [ ] add moral datasets e.g. daily dilemmas, ETHICS, Machiavelli, moral foundations vignettes
Usage
```py
uv sync
uv run src/experiments/ICM.py -- --testbed truthfulQA --alpha 50
uv run nbs/simple_icm.py
```
docs/icm_progress.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 40 KiB

docs/paper.md
+355
View File
File diff suppressed because one or more lines are too long
nbs/simple_icm.py
+143 -57
View File
@@ -14,7 +14,7 @@ import os, sys
from dataclasses import dataclass
import dotenv
from loguru import logger
from openrouter_wrapper.logprobs import openrouter_completion_wlogprobs, get_logprobs_choices # User's wrapper
from openrouter_wrapper.logprobs import openrouter_completion_wlogprobs, get_logprobs_choices, LogprobsNotSupportedError # User's wrapper
from typing import List
dotenv.load_dotenv()
@@ -24,7 +24,6 @@ logger.remove()
logger.add(sys.stderr, format="{time} | {level} | {message}", colorize=True)
# %% [code]
from dataclasses import field
@dataclass
class Config:
@@ -49,14 +48,14 @@ C = Config(
provider_whitelist=[ 'Chutes','Nebius',],
)
C = Config(
model="meta-llama/llama-3.1-70b-instruct",
provider_whitelist=[ 'Cerebras','Nebius',],
)
C = Config(
model="meta-llama/llama-3.1-8b-instruct",
provider_whitelist=[ 'Cerebras','Nebius',],
)
# C = Config(
# model="meta-llama/llama-3.1-70b-instruct",
# provider_whitelist=[ 'Cerebras','Nebius',],
# )
# C = Config(
# model="meta-llama/llama-3.1-8b-instruct",
# provider_whitelist=[ 'Cerebras','Nebius',],
# )
# quick test of logprob
messages = [{'role': 'user',
@@ -96,8 +95,8 @@ for idx, item in enumerate(dataset):
}
data.append(example)
# Limit to small batch for demo (e.g., 64 items)
data = data[:64]
# HACK: Limit to small batch for demo
data = data[:128]
logger.info(f"Loaded {len(data)} examples from TruthfulQA-bool")
# %% [code]
@@ -118,7 +117,10 @@ logger.info("Initialized labels: {}", {k: v['label'] for k, v in demonstrations.
# %% [code]
# Predict label using in-context prompting (placeholder with OpenAI)
prediction_count = 0 # Global counter for first 3 logs
def predict_label(example_uid, current_demos, config=C):
global prediction_count
# Refined few-shot prompt based on get_judge_prompt_fewshot
# Sort by consistency_key for relevance, limit to 8
relevant_demos = sorted(
@@ -139,15 +141,23 @@ def predict_label(example_uid, current_demos, config=C):
try:
# Use wrapper for chat completion (messages format)
messages = [{"role": "user", "content": full_prompt}]
messages = [{"role": "user", "content": full_prompt}, {"role": "assistant", "content": "Judgment: "}]
response = openrouter_completion_wlogprobs(
model_id=config.model,
provider_whitelist=config.provider_whitelist,
messages=messages,
max_tokens=1, # Just for "1" or "0"
max_tokens=3, # Allow for "1" or "0"
temperature=0.0,
top_logprobs=5,
)
# Debug: Log first 3 prompts and responses
if prediction_count < 3:
logger.info(f"Debug Prediction {prediction_count + 1} - UID {example_uid}:")
logger.info(f"Target Prompt: {target_prompt}")
logger.info(f"Response Content: {response['choices'][0]['message']['content']}")
logger.info(f"--- End Debug ---")
prediction_count += 1
# Use wrapper's get_logprobs_choices for score
choice_logp, all_logp = get_logprobs_choices(response, ["1", "0"])
score = choice_logp["1"] - choice_logp["0"]
@@ -156,24 +166,24 @@ def predict_label(example_uid, current_demos, config=C):
# If no logprobs, fallback to text
if response['choices'][0]['logprobs'] is None:
text = response['choices'][0]['message']['content'].strip()
predicted = 1 if "1" in text or "true" in text.lower() else 0
if "0" in text or "false" in text.lower():
predicted = 0
elif "1" in text or "true" in text.lower():
predicted = 1
else:
predicted = np.nan
logger.error(f"Unclear prediction text: {text}")
score = 0.0
logger.warning("No logprobs, using text fallback")
logger.info(f"Prediction for UID {example_uid}: {predicted}, score: {score:.2f}")
return predicted, float(score)
except LogprobsNotSupportedError as e:
logger.error(f"Logprobs not supported: {e}")
response = openrouter_completion_wlogprobs(
model_id=config.model,
messages=[{"role": "user", "content": full_prompt}],
max_tokens=5,
temperature=0.0,
# No logprobs
)
text = response['choices'][0]['message']['content'].strip()
predicted = 1 if "1" in text or "true" in text.lower() else 0
return predicted, 0.0
# except LogprobsNotSupportedError as e:
# logger.error(f"Logprobs not supported: {e}")
# text = response['choices'][0]['message']['content'].strip()
# predicted = 1 if "1" in text or "true" in text.lower() else 0
# return predicted, 0.0
except Exception as e:
logger.error(f"API error: {e}")
return random.choice([0, 1]), 0.0 # Fallback
@@ -184,32 +194,101 @@ test_pred, test_score = predict_label(0, current_labeled, C)
logger.info(f"Test prediction: {test_pred}, score: {test_score}")
# %% [code]
# Simplified inconsistency fix: Ensure per group at most one 'True' for differing keys
def fix_inconsistencies(demos):
# Group by consistency_id
groups = {}
for uid, demo in demos.items():
cid = demo['consistency_id']
if cid not in groups:
groups[cid] = []
groups[cid].append(uid)
# Enhanced inconsistency fix: Multi-iter LLM proposals for contradictions/implications (inspired by ICM_tools.py)
# FIXME: Current basic heuristic; uses LLM for decision prompts on inconsistent pairs, simulates outcomes, iterates to resolve
async def fix_inconsistencies(demos, config=C, max_iters=3):
updated = False
for cid, uids in groups.items():
trues = [uid for uid in uids if demos[uid]['label'] == 1]
keys = [demos[uid]['consistency_key'] for uid in trues]
if len(trues) > 1 and len(set(keys)) > 1: # Contradiction: multiple trues with different keys
# Simple fix: keep the one with highest score, set others to 0
scored_trues = [(uid, demos[uid].get('score', 0)) for uid in trues]
best_uid = max(scored_trues, key=lambda x: x[1])[0]
for uid in trues:
if uid != best_uid:
demos[uid]['label'] = 0
updated = True
for iteration in range(max_iters):
# Group by consistency_id
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)
fixes_made = False
for cid, uids in groups.items():
labeled_uids = [uid for uid in uids if demos[uid]['label'] is not None]
if len(labeled_uids) < 2:
continue # Need at least two labeled for conflict
# Find inconsistent pairs (contradiction or implication)
pairs = []
labels = {uid: demos[uid]['label'] for uid in labeled_uids}
keys = {uid: demos[uid]['consistency_key'] for uid in labeled_uids}
for i, uid1 in enumerate(labeled_uids):
for uid2 in labeled_uids[i+1:]:
label1, label2 = labels[uid1], labels[uid2]
key1, key2 = keys[uid1], keys[uid2]
if key1 != key2 and ((label1 == label2 == 1) or (label1 == label2 == 0 and key1 in ['A>B', 'B>A'])):
pairs.append((uid1, uid2, "contradiction"))
elif key1 == key2 and label1 != label2:
pairs.append((uid1, uid2, "implication"))
if not pairs:
continue
# For each pair, use LLM to propose resolution
for uid1, uid2, pair_type in pairs:
claim1 = demos[uid1]
claim2 = demos[uid2]
# Decision prompt (adapted from get_decision_prompt)
decision_prompt = f"""Resolve inconsistency between two claims:
Claim 1: {claim1['prompt'][:-1]} {claim1['label']}
Claim 2: {claim2['prompt'][:-1]} {claim2['label']}
Type: {pair_type}
If contradiction, decide which to set True (1) and False (0).
If implication, decide if both True or both False.
Respond with 1 if keep Claim1 True and Claim2 False, or 0 otherwise."""
try:
messages = [{"role": "user", "content": decision_prompt}]
response = openrouter_completion_wlogprobs(
model_id=config.model,
provider_whitelist=config.provider_whitelist,
messages=messages,
max_tokens=1,
temperature=0.0,
top_logprobs=5,
)
choice_logp = get_logprobs_choices(response, ["1", "0"])[0]
decision_score = choice_logp["1"] - choice_logp["0"]
decision = 1 if decision_score > 0 else 0
except:
decision = 0 # Fallback; prefer Claim1
# Apply decision
if pair_type == "contradiction":
if decision == 1:
demos[uid1]['label'] = 1
demos[uid2]['label'] = 0
else:
demos[uid1]['label'] = 0
demos[uid2]['label'] = 1
else: # implication
if decision == 1:
demos[uid1]['label'] = 1
demos[uid2]['label'] = 1
else:
demos[uid1]['label'] = 0
demos[uid2]['label'] = 0
fixes_made = True
updated = True
if not fixes_made:
break # No more fixes needed in this iteration
return demos, updated
# Test fix
temp_demos, updated = fix_inconsistencies(demonstrations.copy())
# Test fix (now async, so await)
import asyncio
temp_demos = demonstrations.copy()
temp_demos, updated = asyncio.run(fix_inconsistencies(temp_demos, C))
logger.info(f"Fixed inconsistencies: {updated}")
# %% [code]
@@ -247,7 +326,7 @@ logger.info("Initial energy: {}", compute_energy(demonstrations))
# %% [code]
# Main simulated annealing loop
def run_icm(demonstrations, config=C):
async def run_icm(demonstrations, config=C):
energies = []
accuracies = []
current_labeled = {k: v for k, v in demonstrations.items() if v['label'] is not None}
@@ -297,7 +376,7 @@ def run_icm(demonstrations, config=C):
temp_demos = demonstrations.copy()
temp_demos[example_uid]['label'] = new_label
temp_demos[example_uid]['score'] = score
temp_demos, _ = fix_inconsistencies(temp_demos)
temp_demos, _ = await fix_inconsistencies(temp_demos, config=config)
# Compute new energy
new_energy, new_metrics = compute_energy(temp_demos, config)
@@ -319,16 +398,20 @@ def run_icm(demonstrations, config=C):
if iter % 10 == 0:
logger.info("Progress: Labeled {}, Inconsistents: {}", new_metrics['num_labeled'], new_metrics['num_inconsistent'])
# Final async fix if needed
demonstrations, _ = await fix_inconsistencies(demonstrations, config=config)
return demonstrations, energies, accuracies
# %% [code]
# Run the algorithm
final_demos, energies, accuracies = run_icm(demonstrations, C)
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("Energy: {:.2f}", final_energy)
# TODO show vanilla accuracy here for comparison
logger.info("Accuracy vs vanilla: {:.2f}", final_metrics['accuracy'])
logger.info("Labeled: {}/{}", final_metrics['num_labeled'], len(data))
logger.info("Inconsistencies: {}", final_metrics['num_inconsistent'])
@@ -358,6 +441,7 @@ plt.xlabel('Iteration')
plt.ylabel('Accuracy')
plt.tight_layout()
plt.savefig("icm_progress.png")
plt.show()
# %% [markdown]
@@ -365,7 +449,9 @@ plt.show()
# - [x] Load larger HuggingFace 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).
# - FIXME: Implement async calls for batch efficiency (currently sync).
# - FIXME:
# th logprobs-supported models.
# - [ ] Implement async batch predictions in predict_label (use asyncio.gather for concurrent API calls, inspired by pipeline.py)
# - [ ] Enhance fix_inconsistencies with multi-iter LLM proposals (use LLM decisions for contradictions/implications, from ICM_tools.py)
# - it's marked as "basic" because it directly applies LLM decisions without deeper simulation or multiple proposals.
# - in the original ICM.py the LLM generates multiple resolution proposals for inconsistencies, simulates their outcomes (e.g., by temporarily applying them and evaluating metrics like energy), and iterates to select the best one. https://github.com/Jiaxin-Wen/Unsupervised-Elicitation/blob/master/src/experiments/ICM.py
# - [ ] Add caching for predictions (dict/file-based, like save_to_cache in pipeline.py)
# - [ ] Expand metrics in compute_energy (add label distributions, detailed inconsistent_num, from ICM.py)
src/datatypes/enums.py
-34
View File
@@ -1,34 +0,0 @@
__all__ = ["Language", "PromptType"]
from enum import Enum
class Language(Enum):
PYTHON = ("python", "Python")
CPP = ("cpp", "C++")
def __init__(self, code, text):
self.code = code
self.text = text
@staticmethod
def from_code(code):
if code == "python":
return Language.PYTHON
elif code == "cpp":
return Language.CPP
else:
raise Exception(f"Unknown code langauge: {code}")
class PromptType(Enum):
SOLUTION = "solution_generation"
BLUE_TEAM = "blue_team"
RED_TEAM = "red_team"
EVAL = "eval"
class DifficultyEstimationType(Enum):
PROBLEM_ONLY = "problem_only"
PROBLEM_SOLUTION = "problem_solution"
PROBLEM_SOLUTION_EXECUTION = "problem_solution_execution"
src/experiments/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 unsupervised_elicitation.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)
src/experiments/ICM_tools.py
-238
View File
@@ -1,238 +0,0 @@
import asyncio
import json
import random
from collections import Counter
from copy import deepcopy
import numpy as np
from datasets import load_dataset
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):
return {
"train_predict_distribution": Counter(
[i["label"] for i in train_data.values()]
),
"train_label_distribution": Counter(
[i["vanilla_label"] for i in train_data.values()]
),
"train_accuracy": np.mean(
[i["label"] == i["vanilla_label"] for i in train_data.values()]
),
"train_prob": np.mean(
[
i["score"] if i["label"] == 1 else -i["score"]
for i in train_data.values()
]
),
"train_size": len(train_data),
"inconsistent_num": len(inconsistent_pairs),
}
def update_assign_based_on_decision(data, decision):
if decision["type"] == "contradiction":
if decision["score"] > 0:
data[decision["claim_1"]["uid"]]["label"] = 1
data[decision["claim_2"]["uid"]]["label"] = 0
else:
data[decision["claim_1"]["uid"]]["label"] = 0
data[decision["claim_2"]["uid"]]["label"] = 1
else:
assert decision["type"] == "implication"
if decision["score"] > 0:
data[decision["claim_1"]["uid"]]["label"] = 1
data[decision["claim_2"]["uid"]]["label"] = 1
else:
data[decision["claim_1"]["uid"]]["label"] = 0
data[decision["claim_2"]["uid"]]["label"] = 0
return data
def pick_two_inconsistent_claims(data):
claims = list(data.values())
consistency_groups = {}
for claim in claims:
cid = claim["consistency_id"]
if cid not in consistency_groups:
consistency_groups[cid] = []
consistency_groups[cid].append(claim)
inconsistent_pairs = {}
for group in consistency_groups.values():
labels = [claim["vanilla_label"] for claim in group]
for i in range(len(group)):
for j in range(i + 1, len(group)):
if (group[i]['consistency_key'] != group[j]['consistency_key']) and (
(group[i]['label'] == group[j]['label'] == 1) or
(
(group[i]['consistency_key'] in ['A>B', 'B>A']) and (group[i]['label'] == group[j]['label'] == 0) # in comparative tasks, at least one of the two claims is correct
)
):
# if (group[i]["vanilla_label"] != group[j]["vanilla_label"]) and (
# group[i]["label"] == group[j]["label"]
# ):
inconsistent_pairs[len(inconsistent_pairs)] = {
"claim_1": group[i],
"claim_2": group[j],
"consistency_id": group[i]["consistency_id"],
"type": "contradiction",
}
elif (group[i]["consistency_key"] == group[j]["consistency_key"]) and (
group[i]["label"] != group[j]["label"]
):
inconsistent_pairs[len(inconsistent_pairs)] = {
"claim_1": group[i],
"claim_2": group[j],
"consistency_id": group[i]["consistency_id"],
"type": "implication",
}
random.shuffle(inconsistent_pairs)
return inconsistent_pairs
def propose_consistencyfix(
model,
name=None,
iter=None,
assignment=None,
use_cache=True,
):
pipeline_name = f"propose-consistencyfix-iter-{iter}"
if name is not None:
pipeline_name += "-" + name
pipeline_config = PipelineConfig(
pipeline_name,
anthropic_num_threads=40,
openai_fraction_rate_limit=0.99,
num_problems=None,
use_cache=use_cache,
)
pipeline = Pipeline(pipeline_config)
initial_assign = pipeline.add_load_data_step(
"get_assign", load_assignments, assignment
)
pick_claims = pipeline.add_transformation_step(
"pick_two_inconsistent_claims",
pick_two_inconsistent_claims,
dependencies=[initial_assign],
)
get_decision = pipeline.add_query_step(
"decisions",
model,
get_decision_prompt,
extract_decision_logprobs,
dependencies=[pick_claims],
logprobs=20,
max_tokens=1,
use_cache=use_cache,
)
return pipeline
def run_consistencyfix(
model,
name=None,
use_cache=True,
decision_id=None,
decision=None,
iter=None,
assignment=None,
):
pipeline_name = f"consistencyfix-iter-{iter}"
if decision_id is not None:
pipeline_name += f"-{decision_id}"
if name is not None:
pipeline_name += "-" + name
pipeline_config = PipelineConfig(
pipeline_name,
anthropic_num_threads=40,
openai_fraction_rate_limit=0.99,
num_problems=None,
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
)
pick_claims = pipeline.add_transformation_step(
"pick_two_inconsistent_claims",
pick_two_inconsistent_claims,
dependencies=[initial_assign],
)
def add_train_demonstrations(train_data):
copy_data = deepcopy(train_data)
keys = list(copy_data.keys())
values = list(copy_data.values())
saved_keys = [
"prompt",
"question",
"choice",
"choice_2",
"consistency_id",
"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):
train_data[key]["demonstration"] = {
prev_key: prev_value
for j, (prev_key, prev_value) in enumerate(zip(keys, values))
if j != idx
}
return train_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,
)
eval_preds = pipeline.add_eval_step(
"evaluate",
calculate_accuracy,
dependencies=[get_train_preds, pick_claims],
)
return pipeline
src/experiments/plot.py
-24
View File
@@ -1,24 +0,0 @@
import json
from matplotlib import pyplot as plt
import os
for file in os.listdir("."):
if file.startswith("log_"):
print(file)
with open(file, "r") as f:
data = [json.loads(i) for i in f]
data = data[:60]
x = list(range(8, 8 + len(data)))
y_score = [i['score'] for i in data]
y_acc = [i['acc'] for i in data]
# plt.subplot(1, 2, 1)
plt.plot(x, y_score, label=f'Acc {max(y_acc):.2f}')
plt.xlabel("# Searched Claims")
plt.ylabel("Score")
# plt.subplot(1, 2, 2)
# plt.plot(x, y_acc)
# plt.xlabel("# Searched Claims")
# plt.ylabel("Accuracy")
plt.legend()
plt.tight_layout()
plt.savefig("log.png", dpi=500)
src/llm_api/__init__.py
View File
src/llm_api/base_llm.py
-137
View File
@@ -1,137 +0,0 @@
import json
import logging
from enum import Enum, auto
from typing import Dict, List, Optional, Protocol
import attrs
import numpy as np
# from anthropic import AI_PROMPT, HUMAN_PROMPT
from pydantic import BaseModel
PRINT_COLORS = {"user": "cyan", "system": "magenta", "assistant": "light_green"}
LOGGER = logging.getLogger(__name__)
class PromptConfig(BaseModel):
partials: Dict[str, str] = {}
word_limit: Optional[int] = 100
messages: List[Dict[str, str]] = []
messages1: List[Dict[str, str]] = []
messages2: List[Dict[str, str]] = []
vars: Dict[str, str] = {}
class LanguageModelConfig(BaseModel):
model: str
temperature: float = 0.2
top_p: float = 1.0
max_tokens: Optional[int] = None
max_words: int = 10000
min_words: int = 0
num_candidates_per_completion: int = 1
timeout: int = 120
logit_bias: Optional[dict] = None
class StopReason(Enum):
MAX_TOKENS = auto()
STOP_SEQUENCE = auto()
TOOL_USE = auto()
@classmethod
def factory(cls, stop_reason: str) -> "StopReason":
"""
Parses the openai and anthropic stop reasons into a StopReason enum.
"""
if stop_reason in ["max_tokens", "length"]:
return cls.MAX_TOKENS
elif stop_reason in ["stop_sequence", "stop", "end_turn", "eos"]:
return cls.STOP_SEQUENCE
elif stop_reason in ['tool_use', "tool_calls"]:
return cls.TOOL_USE
raise ValueError(f"Invalid stop reason: {stop_reason}")
def __repr__(self):
return self.name
@attrs.frozen()
class LLMResponse:
model_id: str
completion: str
stop_reason: StopReason = attrs.field(converter=StopReason.factory)
cost: float
duration: Optional[float] = None
api_duration: Optional[float] = None
logprobs: Optional[list[dict[str, float]]] = None
def to_dict(self):
return {
"model_id": self.model_id,
"completion": self.completion,
"stop_reason": self.stop_reason.__repr__(), # Convert to some JSON-serializable format.
"duration": self.duration,
"api_duration": self.api_duration,
"cost": self.cost,
"logprobs": self.logprobs,
}
class ModelAPIProtocol(Protocol):
async def __call__(
self,
model_ids: list[str],
prompt,
print_prompt_and_response: bool,
max_attempts: int,
**kwargs,
) -> list[LLMResponse]:
raise NotImplementedError
def messages_to_single_prompt(messages) -> str:
if (
len(messages) >= 2
and messages[0]["role"] == "system"
and messages[1]["role"] == "user"
):
combined_content = messages[0]["content"] + " " + messages[1]["content"]
messages = [{"role": "user", "content": combined_content}] + messages[2:]
prompt = ""
for message in messages:
role = message["role"]
content = message["content"]
tag = AI_PROMPT if role == "assistant" else HUMAN_PROMPT
prompt += f"{tag} {content}"
if tag != AI_PROMPT:
prompt += f"{AI_PROMPT}"
return prompt.strip()
def convert_to_prob(log_prob: dict, tokens: list) -> tuple[float, float, float]:
logit1 = log_prob.get(tokens[0], None)
logit2 = log_prob.get(tokens[1], None)
if logit1 is None:
rating = -100
LOGGER.warning(
f"Missing token0 {tokens[0]} in log_prob, setting rating to -100.0"
)
else:
rating = logit1
if logit1 is None:
logit1 = -100
if logit2 is None:
logit2 = -100
return rating, logit1, logit2
def add_assistant_message(messages: list[dict], assistant_message: str):
last_role = messages[-1]["role"]
if last_role == "assistant":
messages[-1]["content"] += assistant_message
else:
messages.append({"role": "assistant", "content": assistant_message})
return messages
src/llm_api/llm.py
-274
View File
@@ -1,274 +0,0 @@
import asyncio
import json
import logging
import os
from collections import defaultdict
from itertools import chain
from pathlib import Path
from typing import Callable, Literal, Optional, Union
import attrs
from core.llm_api.base_llm import LLMResponse, ModelAPIProtocol
from core.llm_api.openai_llm import (
BASE_MODELS,
GPT_CHAT_MODELS,
OAIBasePrompt,
OAIChatPrompt,
OpenAIBaseModel,
OpenAIChatModel,
)
from unsupervised_elicitation.utils import load_secrets
LOGGER = logging.getLogger(__name__)
@attrs.define()
class ModelAPI:
openai_fraction_rate_limit: float = attrs.field(
default=0.99, validator=attrs.validators.lt(1)
)
organization: str = "NYU_ORG"
print_prompt_and_response: bool = False
_openai_base: OpenAIBaseModel = attrs.field(init=False)
_openai_base_arg: OpenAIBaseModel = attrs.field(init=False)
_openai_chat: OpenAIChatModel = attrs.field(init=False)
running_cost: float = attrs.field(init=False, default=0)
model_timings: dict[str, list[float]] = attrs.field(init=False, default={})
model_wait_times: dict[str, list[float]] = attrs.field(init=False, default={})
def __attrs_post_init__(self):
secrets = load_secrets()
if self.organization is None:
self.organization = "NYU_ORG"
self._openai_base = OpenAIBaseModel(
frac_rate_limit=self.openai_fraction_rate_limit,
organization=secrets[self.organization],
print_prompt_and_response=self.print_prompt_and_response,
)
self._openai_base_arg = OpenAIBaseModel(
frac_rate_limit=self.openai_fraction_rate_limit,
organization=secrets["ARG_ORG"],
print_prompt_and_response=self.print_prompt_and_response,
)
self._openai_chat = OpenAIChatModel(
frac_rate_limit=self.openai_fraction_rate_limit,
organization=secrets[self.organization],
print_prompt_and_response=self.print_prompt_and_response,
)
Path("./prompt_history").mkdir(exist_ok=True)
@staticmethod
def _load_from_cache(save_file):
if not os.path.exists(save_file):
return None
else:
with open(save_file) as f:
cache = json.load(f)
return cache
async def call_single(
self,
model_ids: Union[str, list[str]],
prompt: Union[list[dict[str, str]], str],
max_tokens: int,
print_prompt_and_response: bool = False,
n: int = 1,
max_attempts_per_api_call: int = 10,
num_candidates_per_completion: int = 1,
# is_valid: Callable[[str], bool] = lambda _: True,
parse_fn=lambda _: True,
insufficient_valids_behaviour: Literal[
"error", "continue", "pad_invalids"
] = "error",
**kwargs,
) -> str:
assert n == 1, f"Expected a single response. {n} responses were requested."
responses = await self(
model_ids,
prompt,
max_tokens,
print_prompt_and_response,
n,
max_attempts_per_api_call,
num_candidates_per_completion,
parse_fn,
insufficient_valids_behaviour,
**kwargs,
)
assert len(responses) == 1, "Expected a single response."
return responses[0].completion
async def __call__(
self,
model_ids: Union[str, list[str]],
prompt: Union[list[dict[str, str]], str],
print_prompt_and_response: bool = False,
n: int = 1,
max_attempts_per_api_call: int = 50,
num_candidates_per_completion: int = 1,
parse_fn=None,
use_cache: bool = True,
file_sem: asyncio.Semaphore = None,
insufficient_valids_behaviour: Literal[
"error", "continue", "pad_invalids"
] = "error",
**kwargs,
) -> list[LLMResponse]:
"""
Make maximally efficient API requests for the specified model(s) and prompt.
Args:
model_ids: The model(s) to call. If multiple models are specified, the output will be sampled from the
cheapest model that has capacity. All models must be from the same class (e.g. OpenAI Base,
OpenAI Chat).
prompt: The prompt to send to the model(s). Type should match what's expected by the model(s).
max_tokens: The maximum number of tokens to request from the API
print_prompt_and_response: Whether to print the prompt and response to stdout.
n: The number of completions to request.
max_attempts_per_api_call: Passed to the underlying API call. If the API call fails (e.g. because the
API is overloaded), it will be retried this many times. If still fails, an exception will be raised.
num_candidates_per_completion: How many candidate completions to generate for each desired completion. n*num_candidates_per_completion completions will be generated, then is_valid is applied as a filter, then the remaining completions are returned up to a maximum of n.
parse_fn: post-processing on the generated response
save_path: cache path
use_cache: whether to load from the cache or overwrite it
"""
assert (
"max_tokens_to_sample" not in kwargs
), "max_tokens_to_sample should be passed in as max_tokens."
if isinstance(model_ids, str):
model_ids = [model_ids]
# # trick to double rate limit for most recent model only
def model_id_to_class(model_id: str) -> ModelAPIProtocol:
if model_id in ["gpt-4-base", "gpt-3.5-turbo-instruct"]:
return (
self._openai_base_arg
) # NYU ARG is only org with access to this model
elif model_id in BASE_MODELS:
return self._openai_base
elif model_id in GPT_CHAT_MODELS or "ft:gpt-3.5-turbo" in model_id:
return self._openai_chat
raise ValueError(f"Invalid model id: {model_id}")
model_classes = [model_id_to_class(model_id) for model_id in model_ids]
# assert model_classes == self._openai_base
# if model_classes == self._openai_base:
# assert "gpt" not in model_ids[0]
# kwargs['api_base'] = "https://5jfmglryfots6s-8000.proxy.runpod.net/v1"
if len(set(str(type(x)) for x in model_classes)) != 1:
raise ValueError("All model ids must be of the same type.")
max_tokens = (
kwargs.get("max_tokens") if kwargs.get("max_tokens") is not None else 2000
)
model_class = model_classes[0]
kwargs["max_tokens"] = max_tokens
# Check if current prompt has already been saved in the save file
# If so, directly return previous result
responses = None
if use_cache and kwargs.get("save_path") is not None:
try:
responses = self._load_from_cache(kwargs.get("save_path"))
except:
logging.error(f"invalid cache data: {kwargs.get('save_path')}")
# After loading cache, we do not directly return previous results,
# but continue running it through parse_fn and re-save it.
# This is because we may frequently update the parse_fn during development
if responses is None:
num_candidates = num_candidates_per_completion * n
responses = await model_class(
model_ids,
prompt,
print_prompt_and_response,
max_attempts_per_api_call,
n=num_candidates,
**kwargs,
)
modified_responses = []
for response in responses:
self.running_cost += response["response"]["cost"]
if kwargs.get("metadata") is not None:
response["metadata"] = kwargs.get("metadata")
if parse_fn is not None:
response = parse_fn(response)
self.model_timings.setdefault(response["response"]["model_id"], []).append(
response["response"]["api_duration"]
)
self.model_wait_times.setdefault(
response["response"]["model_id"], []
).append(
response["response"]["duration"] - response["response"]["api_duration"]
)
modified_responses.append(response)
if kwargs.get("save_path") is not None:
if file_sem is not None:
async with file_sem:
with open(kwargs.get("save_path"), "w") as f:
json.dump(modified_responses, f, indent=2)
else:
with open(kwargs.get("save_path"), "w") as f:
json.dump(modified_responses, f, indent=2)
return modified_responses[:n]
def reset_cost(self):
self.running_cost = 0
async def demo():
model_api = ModelAPI(openai_fraction_rate_limit=0.99)
oai_chat_messages = [
[
{"role": "system", "content": "You are gpt-3.5-turbo."},
{"role": "user", "content": "who are you!"},
],
[
{
"role": "system",
"content": "You are gpt-4",
},
{"role": "user", "content": "who are you!"},
],
]
oai_chat_models = ["gpt-3.5-turbo-16k"]
oai_chat_requests = [
model_api(
oai_chat_models,
prompt=message,
max_tokens=16_000,
n=1,
print_prompt_and_response=False,
)
for message in oai_chat_messages
]
answer = await asyncio.gather(*oai_chat_requests)
for responses in answer:
for i in responses:
print(i.completion)
print("=" * 100)
costs = defaultdict(int)
for responses in answer:
for response in responses:
costs[response.model_id] += response.cost
print("-" * 80)
print("Costs:")
for model_id, cost in costs.items():
print(f"{model_id}: ${cost}")
return answer
if __name__ == "__main__":
asyncio.run(demo())
src/llm_api/openai_llm.py
-630
View File
@@ -1,630 +0,0 @@
# %%
import asyncio
import json
import logging
import os
import random
import time
from datetime import datetime
from itertools import cycle
from traceback import format_exc
from typing import Optional, Union
import attrs
import openai
import requests
import tiktoken
from openai.openai_object import OpenAIObject as OpenAICompletion
from tenacity import retry, stop_after_attempt, wait_fixed
from termcolor import cprint
from core.llm_api.base_llm import (
PRINT_COLORS,
LLMResponse,
ModelAPIProtocol,
)
OAIChatPrompt = list[dict[str, str]]
OAIBasePrompt = Union[str, list[str]]
LOGGER = logging.getLogger(__name__)
def count_tokens(text: str) -> int:
return len(tiktoken.get_encoding("cl100k_base").encode(text))
def price_per_token(model_id: str) -> tuple[float, float]:
"""
Returns the (input token, output token) price for the given model id.
"""
if model_id == "gpt-4-1106-preview":
prices = 0.01, 0.03
elif model_id == "gpt-3.5-turbo-1106":
prices = 0.001, 0.002
elif model_id.startswith("gpt-4"):
prices = 0.03, 0.06
elif model_id.startswith("gpt-4-32k"):
prices = 0.06, 0.12
elif model_id.startswith("gpt-3.5-turbo-16k"):
prices = 0.003, 0.004
elif model_id.startswith("gpt-3.5-turbo"):
prices = 0.0015, 0.002
elif model_id == "davinci-002":
prices = 0.002, 0.002
elif model_id == "babbage-002":
prices = 0.0004, 0.0004
elif model_id == "text-davinci-003" or model_id == "text-davinci-002":
prices = 0.02, 0.02
elif "ft:gpt-3.5-turbo" in model_id:
prices = 0.012, 0.016
elif "llama" in model_id.lower() or "mixtral" in model_id.lower():
prices = 0.0015, 0.002
elif "o1" in model_id.lower():
prices = 0.01, 0.03
else:
prices = 0, 0
# raise ValueError(f"Invalid model id: {model_id}")
return tuple(price / 1000 for price in prices)
@attrs.define()
class Resource:
"""
A resource that is consumed over time and replenished at a constant rate.
"""
refresh_rate: float = (
attrs.field()
) # How many units of the resource are replenished per minute
value: float = attrs.field(init=False)
total: float = 0
throughput: float = 0
last_update_time: float = attrs.field(init=False, factory=time.time)
start_time: float = attrs.field(init=False, factory=time.time)
def __attrs_post_init__(self):
self.value = self.refresh_rate
def _replenish(self):
"""
Updates the value of the resource based on the time since the last update.
"""
curr_time = time.time()
self.value = min(
self.refresh_rate,
self.value + (curr_time - self.last_update_time) * self.refresh_rate / 60,
)
self.last_update_time = curr_time
self.throughput = self.total / (curr_time - self.start_time) * 60
def geq(self, amount: float) -> bool:
self._replenish()
return self.value >= amount
def consume(self, amount: float):
"""
Consumes the given amount of the resource.
"""
assert self.geq(
amount
), f"Resource does not have enough capacity to consume {amount} units"
self.value -= amount
self.total += amount
@attrs.define
class OpenAIModel(ModelAPIProtocol):
frac_rate_limit: float
organization: str
print_prompt_and_response: bool = False
model_ids: set[str] = attrs.field(init=False, default=attrs.Factory(set))
# rate limit
token_capacity: dict[str, Resource] = attrs.field(
init=False, default=attrs.Factory(dict)
)
request_capacity: dict[str, Resource] = attrs.field(
init=False, default=attrs.Factory(dict)
)
lock_add: asyncio.Lock = attrs.field(
init=False, default=attrs.Factory(asyncio.Lock)
)
lock_consume: asyncio.Lock = attrs.field(
init=False, default=attrs.Factory(asyncio.Lock)
)
@staticmethod
def _assert_valid_id(model_id: str):
raise NotImplementedError
@staticmethod
async def _get_dummy_response_header(model_id: str):
raise NotImplementedError
@staticmethod
def _count_prompt_token_capacity(prompt, **kwargs) -> int:
raise NotImplementedError
async def _make_api_call(self, prompt, model_id, **params) -> list[LLMResponse]:
raise NotImplementedError
@staticmethod
def _print_prompt_and_response(prompt, responses):
raise NotImplementedError
@staticmethod
def _create_prompt_history_file(prompt):
filename = f"{datetime.now().strftime('%Y-%m-%d %H:%M:%S.%f')[:-3]}_prompt.txt"
with open(os.path.join("prompt_history", filename), "w") as f:
json_str = json.dumps(prompt, indent=4)
json_str = json_str.replace("\\n", "\n")
f.write(json_str)
return filename
@staticmethod
def _add_response_to_prompt_file(prompt_file, responses):
with open(os.path.join("prompt_history", prompt_file), "a") as f:
f.write("\n\n======RESPONSE======\n\n")
json_str = json.dumps(
[response.to_dict() for response in responses], indent=4
)
json_str = json_str.replace("\\n", "\n")
f.write(json_str)
async def add_model_id(self, model_id: str):
self._assert_valid_id(model_id)
if model_id in self.model_ids:
return
# make dummy request to get token and request capacity
model_metadata = await self._get_dummy_response_header(model_id)
token_capacity = int(model_metadata["x-ratelimit-limit-tokens"])
request_capacity = int(model_metadata["x-ratelimit-limit-requests"])
print(
f"got capacities for model {model_id}: {token_capacity}, {request_capacity}"
)
tokens_consumed = token_capacity - int(
model_metadata["x-ratelimit-remaining-tokens"]
)
requests_consumed = request_capacity - int(
model_metadata["x-ratelimit-remaining-requests"]
)
print(
f"consumed capacities for model {model_id}: {tokens_consumed}, {requests_consumed}"
)
token_cap = token_capacity * self.frac_rate_limit
request_cap = request_capacity * self.frac_rate_limit
if model_id in BASE_MODELS:
token_cap *= (
10000 # openai does not track token limit so we can increase it
)
print(f"setting cap for model {model_id}: {token_cap}, {request_cap}")
self.model_ids.add(model_id)
token_capacity = Resource(token_cap)
request_capacity = Resource(request_cap)
token_capacity.consume(min(token_cap, tokens_consumed))
request_capacity.consume(min(request_cap, requests_consumed))
self.token_capacity[model_id] = token_capacity
self.request_capacity[model_id] = request_capacity
async def __llama_call__(
self,
model_ids: list[str],
prompt,
print_prompt_and_response: bool,
max_attempts: int,
**kwargs,
) -> list[LLMResponse]:
kwargs = {
key: value
for key, value in kwargs.items()
if key not in ("save_path", "metadata")
}
start = time.time()
async def attempt_api_call():
api_base_list = [os.environ['LLAMA_API_BASE']]
kwargs["api_base"] = random.choice(api_base_list)
for model_id in cycle(model_ids):
return await asyncio.wait_for(
self._make_api_call(prompt, model_id, start, **kwargs),
timeout=100, # cloudflare has a 100-second limit for a connection to remain open: https://docs.runpod.io/pods/configuration/expose-ports
)
model_ids.sort(
key=lambda model_id: price_per_token(model_id)[0]
) # Default to cheapest model
model_id = model_ids[0]
prompt = self._process_prompt(prompt)
# prompt_file = self._create_prompt_history_file(prompt)
responses: Optional[list[LLMResponse]] = None
for i in range(max_attempts):
try:
responses = await attempt_api_call()
except Exception as e:
error_info = f"Exception Type: {type(e).__name__}, Error Details: {str(e)}, Traceback: {format_exc()}"
LOGGER.warn(
f"Encountered API error: {error_info}.\nRetrying now. (Attempt {i})"
)
await asyncio.sleep(1.5**i)
else:
break
if responses is None:
raise RuntimeError(
f"Failed to get a response from the API after {max_attempts} attempts."
)
if self.print_prompt_and_response or print_prompt_and_response:
self._print_prompt_and_response(prompt, responses)
end = time.time()
LOGGER.debug(f"Completed call to {model_id} in {end - start}s.")
return [
{"prompt": prompt, "response": response.to_dict()} for response in responses
]
async def __call__(
self,
model_ids: list[str],
prompt,
print_prompt_and_response: bool,
max_attempts: int,
**kwargs,
) -> list[LLMResponse]:
if "gpt" not in model_ids[0]:
return await self.__llama_call__(
model_ids, prompt, print_prompt_and_response, max_attempts, **kwargs
)
kwargs = {
key: value
for key, value in kwargs.items()
if key not in ("save_path", "metadata")
}
start = time.time()
async def attempt_api_call():
for model_id in cycle(model_ids):
async with self.lock_consume:
request_capacity, token_capacity = (
self.request_capacity[model_id],
self.token_capacity[model_id],
)
if request_capacity.geq(1) and token_capacity.geq(token_count):
request_capacity.consume(1)
token_capacity.consume(token_count)
else:
await asyncio.sleep(0.01)
continue # Skip this iteration if the condition isn't met
# Make the API call outside the lock
return await asyncio.wait_for(
self._make_api_call(prompt, model_id, start, **kwargs), timeout=120
)
model_ids.sort(
key=lambda model_id: price_per_token(model_id)[0]
) # Default to cheapest model
async with self.lock_add:
for model_id in model_ids:
await self.add_model_id(model_id)
if "tool" in prompt[0]:
kwargs["tools"] = prompt[0]["tool"]
if "response_format" in prompt[0]:
kwargs['response_format'] = prompt[0]['response_format']
prompt = self._process_prompt(prompt)
token_count = self._count_prompt_token_capacity(prompt, **kwargs)
assert (
max(self.token_capacity[model_id].refresh_rate for model_id in model_ids)
>= token_count
), "Prompt is too long for any model to handle."
# prompt_file = self._create_prompt_history_file(prompt)
responses: Optional[list[LLMResponse]] = None
for i in range(max_attempts):
try:
responses = await attempt_api_call()
except Exception as e:
error_info = f"Exception Type: {type(e).__name__}, Error Details: {str(e)}, Traceback: {format_exc()}"
LOGGER.warn(
f"Encountered API error: {error_info}.\nRetrying now. (Attempt {i})"
)
await asyncio.sleep(1.5**i)
else:
break
if responses is None:
raise RuntimeError(
f"Failed to get a response from the API after {max_attempts} attempts."
)
if self.print_prompt_and_response or print_prompt_and_response:
self._print_prompt_and_response(prompt, responses)
end = time.time()
LOGGER.debug(f"Completed call to {model_id} in {end - start}s.")
return [
{"prompt": prompt, "response": response.to_dict()} for response in responses
]
_GPT_4_MODELS = [
"gpt-4o",
"gpt-4",
"gpt-4-0314",
"gpt-4-0613",
"gpt-4-0125-preview",
"gpt-4-32k",
"gpt-4-32k-0314",
"gpt-4-32k-0613",
"gpt-4-1106-preview",
"gpt-4-turbo",
"gpt-4-turbo-preview",
"gpt-4-turbo-2024-04-09",
"gpt-4o-mini",
"gpt-4o-mini-2024-07-18",
"gpt-4o-2024-11-20",
"o1-preview-2024-09-12",
"o1-mini-2024-09-12",
"deepseek/deepseek-chat",
"meta-llama/llama-3.2-3b-instruct",
"meta-llama/llama-3.2-1b-instruct",
"meta-llama/llama-3.3-70b-instruct",
"mistralai/mistral-7b-instruct",
"meta-llama/llama-3-8b-instruct",
"allenai/olmo-7b-instruct",
"01-ai/yi-large",
"meta-llama/llama-2-70b-chat",
"meta-llama/llama-3.1-8b-instruct",
"meta-llama/llama-3.1-70b-instruct",
"meta-llama/llama-3.1-405b-instruct",
"qwen/qwen-2.5-7b-instruct",
"openai/gpt-4o",
"openchat/openchat-7b",
"ai21/jamba-instruct",
"neversleep/llama-3.1-lumimaid-8b",
"mistralai/mixtral-8x7b-instruct:nitro",
"deepseek/deepseek-r1",
"deepseek/deepseek-r1-distill-llama-70b",
"minimax/minimax-01",
"microsoft/phi-4",
"qwen/qvq-72b-preview",
]
_GPT_TURBO_MODELS = [
"gpt-3.5-turbo",
"gpt-3.5-turbo-0613",
"gpt-3.5-turbo-16k",
"gpt-3.5-turbo-16k-0613",
"gpt-3.5-turbo-1106",
"gpt-3.5-turbo-0125",
]
GPT_CHAT_MODELS = set(_GPT_4_MODELS + _GPT_TURBO_MODELS)
class OpenAIChatModel(OpenAIModel):
def _process_prompt(self, prompt: OAIChatPrompt) -> OAIChatPrompt:
return prompt
def _assert_valid_id(self, model_id: str):
if "ft:" in model_id:
model_id = model_id.split(":")[1]
assert model_id in GPT_CHAT_MODELS, f"Invalid model id: {model_id}"
@retry(stop=stop_after_attempt(8), wait=wait_fixed(2))
async def _get_dummy_response_header(self, model_id: str):
url = "https://api.openai.com/v1/chat/completions"
headers = {
"Content-Type": "application/json",
"Authorization": f"Bearer {openai.api_key}",
"OpenAI-Organization": self.organization,
}
data = {
"model": model_id,
"messages": [{"role": "user", "content": "Say 1"}],
}
response = requests.post(url, headers=headers, json=data)
if "x-ratelimit-limit-tokens" not in response.headers:
raise RuntimeError("Failed to get dummy response header")
return response.headers
@staticmethod
def _count_prompt_token_capacity(prompt: OAIChatPrompt, **kwargs) -> int:
# The magic formula is: .25 * (total number of characters) + (number of messages) + (max_tokens, or 15 if not specified)
BUFFER = 5 # A bit of buffer for some error margin
MIN_NUM_TOKENS = 20
num_tokens = 0
for message in prompt:
num_tokens += 1
num_tokens += len(message["content"]) / 4
return max(
MIN_NUM_TOKENS,
int(num_tokens + BUFFER)
+ kwargs.get("n", 1) * kwargs.get("max_tokens", 15),
)
def convert_top_logprobs(self, data):
# Initialize the new structure with only top_logprobs
top_logprobs = []
for item in data["content"]:
# Prepare a dictionary for top_logprobs
top_logprob_dict = {}
for top_logprob in item["top_logprobs"]:
top_logprob_dict[top_logprob["token"]] = top_logprob["logprob"]
top_logprobs.append(top_logprob_dict)
return top_logprobs
async def _make_api_call(
self, prompt: OAIChatPrompt, model_id, start_time, **params
) -> list[LLMResponse]:
LOGGER.debug(f"Making {model_id} call with {self.organization}")
if params.get("logprobs", None):
params["top_logprobs"] = params["logprobs"]
params["logprobs"] = True
api_start = time.time()
api_response: OpenAICompletion = await openai.ChatCompletion.acreate(messages=prompt, model=model_id, organization=self.organization, **params) # type: ignore
api_duration = time.time() - api_start
duration = time.time() - start_time
context_token_cost, completion_token_cost = price_per_token(model_id)
context_cost = api_response.usage.prompt_tokens * context_token_cost
completion_cost = api_response.usage.completion_tokens * completion_token_cost
return [
LLMResponse(
model_id=model_id,
completion=choice.message.content
if "tools" not in params
else choice.message.tool_calls[0]["function"]["arguments"],
stop_reason=choice.finish_reason,
api_duration=api_duration,
duration=duration,
cost=context_cost + completion_cost,
logprobs=self.convert_top_logprobs(choice.logprobs)
if choice.logprobs is not None
else None,
)
for choice in api_response.choices
]
@staticmethod
def _print_prompt_and_response(
prompts: OAIChatPrompt, responses: list[LLMResponse]
):
for prompt in prompts:
role, text = prompt["role"], prompt["content"]
cprint(f"=={role.upper()}:", "white")
cprint(text, PRINT_COLORS[role])
for i, response in enumerate(responses):
if len(responses) > 1:
cprint(f"==RESPONSE {i + 1} ({response.model_id}):", "white")
cprint(response.completion, PRINT_COLORS["assistant"], attrs=["bold"])
print()
BASE_MODELS = {
"meta-llama/Llama-3.1-8B",
"meta-llama/Llama-3.1-70B",
}
class OpenAIBaseModel(OpenAIModel):
def _process_prompt(
self, prompt: Union[OAIBasePrompt, OAIChatPrompt]
) -> OAIBasePrompt:
if isinstance(prompt, list) and isinstance(prompt[0], dict):
return messages_to_single_prompt(prompt)
return prompt
def _assert_valid_id(self, model_id: str):
assert model_id in BASE_MODELS, f"Invalid model id: {model_id}"
@retry(stop=stop_after_attempt(8), wait=wait_fixed(2))
async def _get_dummy_response_header(self, model_id: str):
url = "https://api.openai.com/v1/completions"
headers = {
"Content-Type": "application/json",
"Authorization": f"Bearer {openai.api_key}",
"OpenAI-Organization": self.organization,
}
data = {"model": model_id, "prompt": "a", "max_tokens": 1}
response = requests.post(url, headers=headers, json=data)
if "gpt" in model_id and "x-ratelimit-limit-tokens" not in response.headers:
raise RuntimeError("Failed to get dummy response header")
return response.headers
@staticmethod
def _count_prompt_token_capacity(prompt: OAIBasePrompt, **kwargs) -> int:
max_tokens = kwargs.get("max_tokens", 15)
n = kwargs.get("n", 1)
completion_tokens = n * max_tokens
tokenizer = tiktoken.get_encoding("cl100k_base")
if isinstance(prompt, str):
prompt_tokens = len(tokenizer.encode(prompt))
return prompt_tokens + completion_tokens
else:
prompt_tokens = sum(len(tokenizer.encode(p)) for p in prompt)
return prompt_tokens + completion_tokens
async def _make_api_call(
self, prompt: OAIBasePrompt, model_id, start_time, **params
) -> list[LLMResponse]:
LOGGER.debug(f"Making {model_id} call with {self.organization}")
api_start = time.time()
api_response: OpenAICompletion = await openai.Completion.acreate(prompt=prompt, model=model_id, organization=self.organization, **params) # type: ignore
api_duration = time.time() - api_start
duration = time.time() - start_time
if "gpt" not in model_id:
return [
LLMResponse(
model_id=model_id,
completion=choice.text,
stop_reason=choice.finish_reason,
api_duration=api_duration,
duration=duration,
cost=0,
logprobs=choice.logprobs.top_logprobs
if choice.logprobs is not None
else None,
)
for choice in api_response.choices
]
else:
context_token_cost, completion_token_cost = price_per_token(model_id)
context_cost = api_response.usage.prompt_tokens * context_token_cost
return [
LLMResponse(
model_id=model_id,
completion=choice.text,
stop_reason=choice.finish_reason,
api_duration=api_duration,
duration=duration,
cost=context_cost / len(api_response.choices)
+ count_tokens(choice.message.content) * completion_token_cost,
logprobs=choice.logprobs.top_logprobs
if choice.logprobs is not None
else None,
)
for choice in api_response.choices
]
@staticmethod
def _print_prompt_and_response(prompt: OAIBasePrompt, responses: list[LLMResponse]):
prompt_list = prompt if isinstance(prompt, list) else [prompt]
responses_per_prompt = len(responses) // len(prompt_list)
responses_list = [
responses[i : i + responses_per_prompt]
for i in range(0, len(responses), responses_per_prompt)
]
for i, (prompt, response_list) in enumerate(zip(prompt_list, responses_list)):
if len(prompt_list) > 1:
cprint(f"==PROMPT {i + 1}", "white")
if len(response_list) == 1:
cprint(f"=={response_list[0].model_id}", "white")
cprint(prompt, PRINT_COLORS["user"], end="")
cprint(
response_list[0].completion,
PRINT_COLORS["assistant"],
attrs=["bold"],
)
else:
cprint(prompt, PRINT_COLORS["user"])
for j, response in enumerate(response_list):
cprint(f"==RESPONSE {j + 1} ({response.model_id}):", "white")
cprint(
response.completion, PRINT_COLORS["assistant"], attrs=["bold"]
)
print()
# %%
src/model_querying/__init__.py
View File
src/model_querying/prompt_creation.py
-45
View File
@@ -1,45 +0,0 @@
__all__ = ["get_apps_solution_prompt", "get_two_shot_prompt", "get_backdoor_prompt"]
import json
import random
from src.datatypes.enums import Language, PromptType
from src.tools.dataloaders import load_prompts
from src.tools.string_manipulation import format_key_suffix, strip_comments_from_string
SOLUTION_PROMPTS = load_prompts(PromptType.SOLUTION)
RED_TEAM_PROMPTS = load_prompts(PromptType.RED_TEAM)
BLUE_TEAM_PROMPTS = load_prompts(PromptType.BLUE_TEAM)
EVAL_PROMPTS = load_prompts(PromptType.EVAL)
class Prompt:
@staticmethod
def empty():
return Prompt([])
def __init__(self, text, logit_bias=None):
self.text = text
self.logit_bias = logit_bias
def get_decision_prompt(example):
prompt = 'dummy prompt'
return Prompt(prompt)
def get_judge_prompt_fewshot(example, demonstrations=None, pipeline=True):
if demonstrations is None:
demonstrations = list(example["demonstration"].values())
prompt = ""
for i in demonstrations:
prompt += i['prompt']
prompt += "True" if i["label"] else "False"
prompt += "\n\n"
prompt += example['prompt']
if pipeline:
return Prompt(prompt)
else:
return prompt
src/model_querying/solution_extraction.py
-55
View File
@@ -1,55 +0,0 @@
import json
import logging
import math
from copy import copy
logger = logging.getLogger(__name__)
def get_yes_no(x):
x = x.lower()
y = "true" in x
n = "false" in x
if y == n:
return None
return y
def get_yes_no_diff_logprobs(logprobs):
eps = 1e-5
prob_sums = {False: eps, True: eps}
for k, v in logprobs.items():
o = get_yes_no(k)
if o is None:
continue
prob_sums[o] += math.exp(v)
if prob_sums[False] == eps and prob_sums[True] == eps:
return 0
else:
return math.log(prob_sums[True]) - math.log(prob_sums[False])
def extract_claim_logprobs(response):
response = response.copy()
try:
logprobs = response["response"]["logprobs"][0]
response[f"score"] = get_yes_no_diff_logprobs(logprobs)
except Exception as e:
logger.info(
f"Problem {response['metadata']['uid']}: Error extracting judgment: {repr(e)}"
)
response["score"] = 0
return response
def extract_decision_logprobs(response):
response = response.copy()
try:
logprobs = response["response"]["logprobs"][0]
response[f"score"] = get_yes_no_diff_logprobs(logprobs)
except Exception as e:
logger.info(
f"Problem {response['metadata']['uid']}: Error extracting decision: {repr(e)}"
)
response["score"] = 0
return response
src/pipeline/README.md
-27
View File
@@ -1,27 +0,0 @@
How to use the pipeline:
First, outline the graph you would like to execute, including the following:
* Data Loading
* Model queries
* Code Execution Eval
* Transformations
* Monitoring
Next, convert each of the nodes in that graph into the corresponding helper function:
* add_load_data_step
* add_query_step
* add_code_evaluation_step
* add_transformation_step
* add_monitoring_step
Each of these takes different parameters that you can see in the method signatures. The important ones to know are these:
LoadData takes either a data-loading function and a location, or it takes a dataset
Queries take a prompt function that they pass the incoming data into to create the associated prompt, and a parse function that they use to parse the LLM response
Code Evals take an executor function that executes all of the code in the Solution objects on the associated test cases.
Transforms take arbitrary functions that they apply to the data as a whole.
Monitoring steps take arbitrary monitoring steps. I may eventually enforce that all pipelines end in one of these because it's really what we care about.
Finally, put the dependencies of each step into their dependencies parameter. This is how execution order is determined and how data flows between steps. If you rely on more than one step, the data will be passed to ordered args in the same order as the list of dependencies.
You will also need to include a PipelineConfig parameter that contains metadata around how many concurrents to use and similar.
Once you have a pipeline definition, call the pipeline.run() function on it to execute the graph. This method returns the Pipeline.Results object back, which holds the output of each step in a dictionary.
src/pipeline/pipeline.py
-314
View File
@@ -1,314 +0,0 @@
__all__ = ["PipelineConfig", "Pipeline"]
import asyncio
import logging
from collections import deque
from tqdm.auto import tqdm
from core.llm_api.llm import ModelAPI
from src.datatypes.enums import Language
from src.runners.query_model import QueryConfigBuilder, query_model
from src.tools.dataloaders import read_from_cache, save_to_cache
from src.tools.path_utils import get_root_directory
logger = logging.getLogger(__name__)
def in_notebook():
try:
from IPython import get_ipython
if get_ipython() is None or "IPKernelApp" not in get_ipython().config:
return False
except ImportError:
return False
return True
class Task:
def __init__(self, name, func, use_cache, dependencies=[]):
self.name = name
self.func = func
self.use_cache = use_cache
self.index = None
self.dependencies = dependencies
self.dependents = []
self.result = None
for dep in dependencies:
dep.dependents.append(self)
async def execute(self, results):
if self.result is None:
dep_results = [results[dep.name] for dep in self.dependencies]
if asyncio.iscoroutinefunction(self.func):
self.result = await self.func(
*dep_results, use_cache=self.use_cache, index=self.index
)
else:
self.result = self.func(
*dep_results, use_cache=self.use_cache, index=self.index
)
return self.result
class PipelineConfig:
def __init__(
self,
name,
anthropic_num_threads=2,
openai_fraction_rate_limit=0.99,
use_cache=True,
language=Language.PYTHON,
num_problems=None,
problem_ids=None,
num_open_files=1000000,
organization="NYU_ORG",
print_prompt_and_response=False,
api_base=None,
):
self.name = name
self.anthropic_num_threads = anthropic_num_threads
self.openai_fraction_rate_limit = openai_fraction_rate_limit
self.organization = organization
self.print_prompt_and_response = print_prompt_and_response
self.use_cache = use_cache
self.language = language
self.num_problems = num_problems
self.problem_ids = problem_ids
self.num_open_files = num_open_files
self.api_base = api_base
self.play_sound = in_notebook()
class Pipeline:
def __init__(self, config):
self.config = config
self.steps = []
self.step_names = set()
self.results = {}
self.model_api = ModelAPI(
self.config.anthropic_num_threads,
self.config.openai_fraction_rate_limit,
self.config.organization,
self.config.print_prompt_and_response,
)
self.file_sem = asyncio.BoundedSemaphore(self.config.num_open_files)
self.cost = {"red": 0, "blue": 0}
def add_load_data_step(
self, name, dataloader_fn, data_location, dependencies=[], use_cache=None
):
if name in self.step_names:
raise ValueError(f"Step name {name} already exists")
self.step_names.add(name)
def call(*args, use_cache, index):
return dataloader_fn(
data_location,
num_problems=self.config.num_problems,
problem_ids=self.config.problem_ids,
)
task = Task(name, call, use_cache, dependencies)
self.steps.append(task)
return task
def add_query_step(
self,
name,
model,
prompt_fn,
parse_fn,
dependencies=[],
use_cache=None,
temperature=None,
logprobs=None,
team=None,
max_tokens=4096,
bon=1,
):
if name in self.step_names:
raise ValueError(f"Step name {name} already exists")
self.step_names.add(name)
query_config_builder = (
QueryConfigBuilder()
.with_model_to_test(model)
.with_prompt_fn(lambda x: prompt_fn(x))
.with_parse_fn(lambda x: parse_fn(x))
.with_num_problems(self.config.num_problems)
.with_max_tokens(max_tokens)
.with_temperature(temperature)
.with_logprobs(logprobs)
.with_bon(bon)
)
async def call(data, use_cache, index):
response_dict = await query_model(
self.model_api,
self.file_sem,
query_config_builder.with_experiment_name(
f"{self.config.name}/{index:02d}-{name}"
)
.with_use_cache(use_cache)
.with_data(data)
.build(),
)
self.add_cost_data(team, response_dict)
return response_dict
step = Task(name, call, use_cache, dependencies)
self.steps.append(step)
return step
def add_transformation_step(
self,
name,
transformation_fn,
dependencies=[],
use_cache=None,
strong_model=None,
weak_model=None,
read_cache=False,
):
if name in self.step_names:
raise ValueError(f"Step name {name} already exists")
self.step_names.add(name)
async def call(*args, use_cache, index):
incoming_problem_ids = set().union(*[arg.keys() for arg in args])
if use_cache and read_cache:
logger.debug(
f"Reading from cache for transformation: {self.config.name}/{index:02d}-{name}/{strong_model}{'+' if strong_model and weak_model else ''}{weak_model}"
)
data, cached_problem_ids = read_from_cache(
f"{self.config.name}/{index:02d}-{name}/{strong_model}{'+' if strong_model and weak_model else ''}{weak_model}"
)
if incoming_problem_ids.issubset(set(cached_problem_ids)):
return {k: v for k, v in data.items() if k in incoming_problem_ids}
if asyncio.iscoroutinefunction(transformation_fn):
output = await transformation_fn(*args)
else:
output = transformation_fn(*args)
async with self.file_sem:
save_to_cache(
output,
f"{self.config.name}/{index:02d}-{name}/{strong_model}{'+' if strong_model and weak_model else ''}{weak_model}",
delete_existing=read_cache,
incoming_problem_ids=incoming_problem_ids,
)
return output
step = Task(name, call, use_cache, dependencies)
self.steps.append(step)
return step
def add_eval_step(
self,
name,
eval_fn,
dependencies=[],
strong_model=None,
weak_model=None,
):
if name in self.step_names:
raise ValueError(f"Step name {name} already exists")
self.step_names.add(name)
async def call(*args, use_cache, index):
output = eval_fn(*args)
cache_obj = {"summary": output}
async with self.file_sem:
save_to_cache(
cache_obj,
f"{self.config.name}/{index:02d}-{name}/{strong_model}{'+' if strong_model and weak_model else ''}{weak_model}",
)
return output
step = Task(name, call, None, dependencies)
self.steps.append(step)
return step
def topological_sort_tasks(self, tasks):
in_degree = {task: len(task.dependencies) for task in tasks}
queue = deque([task for task in tasks if in_degree[task] == 0])
sorted_tasks = []
task_order = {task: i for i, task in enumerate(tasks)}
while queue:
task = queue.popleft()
sorted_tasks.append(task)
for dependent in task.dependents:
in_degree[dependent] -= 1
if in_degree[dependent] == 0:
queue.append(dependent)
queue = deque(sorted(queue, key=lambda t: task_order[t]))
for i, task in enumerate(sorted_tasks):
task.index = i
return sorted_tasks
def add_cost_data(self, team, response_dict):
cost = sum(
[response["response"]["cost"] for response in response_dict.values()]
)
if team is not None:
if team not in self.cost:
self.cost[team] = 0
self.cost[team] += cost
overall_team = team.split("_")[0]
if overall_team != team:
self.cost[overall_team] += cost
def set_use_cache(self, tasks):
# This is called after tasks.sort, so we are guaranteed to process all
# dependencies before each task itself.
for task in tasks:
if not self.config.use_cache:
task.use_cache = False
continue
if task.use_cache is None:
task.use_cache = True
for dep in task.dependencies:
if not dep.use_cache:
task.use_cache = False
def speak(self, message):
if self.config.play_sound:
from IPython.display import Javascript, display
display(
Javascript(
f"""
if(window.speechSynthesis) {{
var synth = window.speechSynthesis;
synth.speak(new window.SpeechSynthesisUtterance('{message}'));
}}
"""
)
)
async def run(self):
steps = self.topological_sort_tasks(self.steps)
self.set_use_cache(steps)
for task in steps:
logger.info(
f"Starting step {task.index}: {task.name} - Using cache: {task.use_cache}"
)
try:
self.results[task.name] = await task.execute(self.results)
except Exception as e:
logger.error(f"Error in step {task.index}: {task.name}")
logger.error(e)
self.speak("Pipeline failed sad face")
raise e
logger.info(f"Finished step {task.index}: {task.name}")
self.speak("Jobs done")
logger.info("Run complete!! Nice!! 🚀🚀")
return self.results
src/runners/__init__.py
View File
src/runners/evaluate_code.py
-191
View File
@@ -1,191 +0,0 @@
__all__ = [
"EvalConfig",
"EvalConfigBuilder",
"evaluate_solutions",
"examine_solution",
"print_eval",
]
import json
import os
from src.code_evaluation.test_results import Solution
import src.tools.path_utils as path_utils
DEFAULT_RESULTS_DIR = path_utils.get_default_results_directory()
class EvalConfig:
def __init__(
self,
experiment_name,
model_to_test,
executor_fn,
language,
use_cache=True,
data=None,
dataloader_fn=None,
data_location=None,
):
self.experiment_name = experiment_name
self.model_to_test = model_to_test
if dataloader_fn is not None:
self.data = dataloader_fn(data_location)
elif isinstance(data, str):
with open(json.load(data), "r") as f:
self.data = json.load(f)
else:
self.data = data
self.executor_fn = executor_fn
self.language = language
self.use_cache = use_cache
def __str__(self):
pass
def __repr__(self):
return self.__str__()
class EvalConfigBuilder:
def __init__(self):
self.experiment_name = None
self.model_to_test = None
self.executor_fn = None
self.language = None
self.data = None
self.use_cache = None
self.dataloader_fn = None
self.data_location = None
def with_experiment_name(self, experiment_name):
self.experiment_name = experiment_name
return self
def with_executor_fn(self, executor_fn):
self.executor_fn = executor_fn
return self
def with_language(self, language):
self.language = language
return self
def with_model_to_test(self, model_to_test):
self.model_to_test = model_to_test
return self
def with_use_cache(self, use_cache):
self.use_cache = use_cache
return self
def with_data(self, data):
self.data = data
return self
def with_dataloader_fn(self, dataloader_fn):
self.dataloader_fn = dataloader_fn
return self
def with_data_location(self, data_location):
self.data_location = data_location
return self
def build(self):
assert self.experiment_name is not None, "Experiment name must be set"
assert self.executor_fn is not None, "Executor function must be set"
assert self.language is not None, "Language must be set"
assert (self.data is not None) or (
self.dataloader_fn is not None and self.data_location is not None
), "Data must be set"
return EvalConfig(
self.experiment_name,
self.model_to_test,
self.executor_fn,
self.language,
self.use_cache,
self.data,
self.dataloader_fn,
self.data_location,
)
def evaluate_solutions(eval_config):
if isinstance(eval_config.data, list):
eval_config.data = {
index: response[0] for index, response in enumerate(eval_config.data)
}
eval_data = []
for problem_id, response in eval_config.data.items():
if response == {}:
item = Solution.no_solution(problem_id)
else:
item = Solution.from_response(problem_id, response, eval_config.language)
eval_data.append(item)
results_dir = DEFAULT_RESULTS_DIR
save_dir = results_dir / eval_config.experiment_name / eval_config.model_to_test
os.makedirs(save_dir, exist_ok=True)
# load caching results
if eval_config.use_cache and save_dir is not None:
for idx, solution in enumerate(eval_data):
if solution.correct is not None: # already have runtime eval results
continue
save_path = f"{save_dir}/{solution.question_id}.json"
if os.path.exists(save_path):
eval_data[idx] = Solution.from_cache_file(save_path)
executor_results = eval_config.executor_fn(eval_data)
for problem_id, problem_data in eval_config.data.items():
if problem_id not in executor_results:
executor_results[problem_id] = {}
for field, value in problem_data.items():
if field not in executor_results[problem_id]:
executor_results[problem_id][field] = value
# save caching results
if save_dir is not None:
for question_id, solution in executor_results.items():
save_path = f"{save_dir}/{question_id}.json"
with open(save_path, "w") as f:
json.dump(solution, f, indent=2)
return executor_results
def examine_solution(solutions, index):
print(
f"Difficulty:\n{solutions[index][0]['metadata']['difficulty']}\n----------------------------"
)
print(solutions[index][0]["metadata"]["question"])
print(solutions[index][0]["solution"])
print("test cases:")
for test in solutions[index][0]["metadata"]["test_cases"]:
print(f"{test['input']}{test['output']}")
def print_eval(results):
correct_tests = 0
total_tests = 0
correct_problems = 0
total_problems = len(results.keys())
for problem_id, problem_result in results.items():
correct_tests_local = sum(
[
1
for test_result in problem_result["test_cases"]
if test_result["correct"]
]
)
total_tests_local = len(problem_result["test_cases"])
# print(f"Problem ID: {problem_id}\nTest Results:\n\tCorrect: {correct_tests_local}\n\tTotal: {total_tests_local}\n\tAccuracy: {(correct_tests_local * 100.)/total_tests_local}\nOverall Correct: {problem_result.correct}")
correct_tests += correct_tests_local
total_tests += total_tests_local
if problem_result["correct"]:
correct_problems += 1
print(
f"Number of Problems: {total_problems}\nNumber Correct: {correct_problems}\nAccuracy: {(correct_problems * 100.)/total_problems}\nNumber of Tests: {total_tests}\nNumber Correct: {correct_tests}\nAccuracy: {(correct_tests * 100.)/total_tests}"
)
src/runners/query_model.py
-276
View File
@@ -1,276 +0,0 @@
__all__ = ["QueryConfig", "QueryConfigBuilder", "query_model"]
import asyncio
import os
import tiktoken
import src.tools.path_utils as path_utils
ROOT_DIR = path_utils.get_root_directory()
DEFAULT_RESULTS_DIR = path_utils.get_default_results_directory()
class QueryConfig:
def __init__(
self,
experiment_name,
model_to_test,
dataloader_fn,
data_location,
data,
prompt_fn,
parse_fn=None,
use_cache=False,
num_problems=None,
max_tokens=4096,
results_dir=None,
temperature=None,
logprobs=None,
bon=1,
):
assert isinstance(model_to_test, str)
self.experiment_name = experiment_name
self.model_to_test = model_to_test
self.dataloader_fn = dataloader_fn
self.data_location = data_location
self.data = data
self.prompt_fn = prompt_fn
self.use_cache = use_cache
self.parse_fn = parse_fn
self.num_problems = num_problems
self.max_tokens = max_tokens
self.results_dir = results_dir
self.temperature = temperature if temperature is not None else 0.0
self.logprobs = logprobs
self.bon = bon
def get_data(self):
assert self.data is not None
return self.data
def __str__(self):
return (
f"QueryConfig("
f"experiment_name={self.experiment_name}, "
f"model_to_test={self.model_to_test}, "
f"dataloader_fn={self.dataloader_fn}, "
f"data_location={self.data_location}, "
f"data={self.data}, "
f"prompt_fn={self.prompt_fn}, "
f"use_cache={self.use_cache}, "
f"num_problems={self.num_problems}, "
f"max_tokens={self.max_tokens}, "
f"results_dir={self.results_dir}, "
f"temperature={self.temperature}, "
f"logprobs={self.logprobs}"
)
def __repr__(self):
return self.__str__()
class QueryConfigBuilder:
def __init__(self):
self.experiment_name = None
self.model_to_test = None
self.dataloader_fn = None
self.data_location = None
self.data = None
self.prompt_fn = None
self.parse_fn = None
self.use_cache = False
self.num_problems = None
self.max_tokens = 4096
self.results_dir = None
self.temperature = 0.0
self.logprobs = None
self.bon = 1
def with_experiment_name(self, experiment_name):
self.experiment_name = experiment_name
return self
def with_bon(self, bon):
self.bon = bon
return self
def with_model_to_test(self, model_to_test):
assert isinstance(model_to_test, str)
self.model_to_test = model_to_test
return self
def with_dataloader_fn(self, dataloader_fn):
self.dataloader_fn = dataloader_fn
return self
def with_data_location(self, data_location):
self.data_location = data_location
return self
def with_data(self, data):
self.data = data
return self
def with_prompt_fn(self, prompt_fn):
self.prompt_fn = prompt_fn
return self
def with_parse_fn(self, parse_fn):
self.parse_fn = parse_fn
return self
def with_use_cache(self, use_cache):
self.use_cache = use_cache
return self
def with_num_problems(self, num_problems):
self.num_problems = num_problems
return self
def with_max_tokens(self, max_tokens):
self.max_tokens = max_tokens
return self
def with_results_dir(self, results_dir):
self.results_dir = results_dir
return self
def with_temperature(self, temperature):
self.temperature = temperature
return self
def with_logprobs(self, logprobs):
self.logprobs = logprobs
if logprobs is not None:
assert "claude" not in self.model_to_test
return self
def build(self):
assert self.experiment_name is not None, "Experiment name must be set"
assert self.model_to_test is not None, "Model to test must be set"
assert self.prompt_fn is not None, "Prompt function must be set"
assert (self.data is not None) or (
self.dataloader_fn is not None and self.data_location is not None
), "Data must be set"
assert (self.data is None) or (
self.dataloader_fn is None and self.data_location is None
), "Data and dataloader_fn/data_location cannot both be set"
return QueryConfig(
self.experiment_name,
self.model_to_test,
self.dataloader_fn,
self.data_location,
self.data,
self.prompt_fn,
self.parse_fn,
self.use_cache,
self.num_problems,
self.max_tokens,
self.results_dir,
self.temperature,
self.logprobs,
self.bon,
)
def _get_prompts(problems, prompt_fn):
prompts = {}
for problem_id, problem in problems.items():
prompt = prompt_fn(problem)
if prompt.text:
prompts[problem_id] = prompt
return prompts
def get_save_dir(query_config):
results_dir = (
ROOT_DIR / query_config.results_dir
if query_config.results_dir is not None
else DEFAULT_RESULTS_DIR
)
save_dir = results_dir / query_config.experiment_name / query_config.model_to_test
os.makedirs(save_dir, exist_ok=True)
return save_dir
def move_data_into_metadata(data):
for data_id, value in data.items():
filtered_val = {
k: v
for k, v in value.items()
if k not in ("metadata", "prompt", "response")
}
metadata = value.get("metadata", {})
new_metadata = metadata | filtered_val
data[data_id]["metadata"] = new_metadata
def format_response(data, model_responses_map, query_config):
model_responses_flattened = {}
for data_id, response in model_responses_map.items():
if query_config.bon == 1:
model_responses_flattened[f"{data_id}"] = data[data_id] | response[0]
continue
for resp_id, resp in enumerate(response):
model_responses_flattened[f"{data_id}-{resp_id}"] = data[data_id] | resp
return model_responses_flattened
def tokenize_logit_bias(logit_bias, model):
tokenizer = tiktoken.encoding_for_model(model)
tokenized_bias = {}
for k, v in logit_bias.items():
tokenized = tokenizer.encode(k)
assert len(tokenized) == 1, f"Tokenized bias key {k} is not a single token"
tokenized_bias[tokenized[0]] = v
return tokenized_bias
async def query_model(model_api, file_sem, query_config):
data = query_config.get_data()
prompts = _get_prompts(data, query_config.prompt_fn)
save_dir = get_save_dir(query_config)
move_data_into_metadata(data)
model_requests = [
model_api(
query_config.model_to_test,
prompts[data_id].text,
max_tokens=query_config.max_tokens,
temperature=query_config.temperature,
n=query_config.bon,
top_p=1.0,
logprobs=query_config.logprobs,
use_cache=query_config.use_cache,
metadata=data[data_id]["metadata"],
parse_fn=query_config.parse_fn,
save_path=f"{save_dir}/{data_id}.json",
file_sem=file_sem,
**(
{
"logit_bias": tokenize_logit_bias(
prompts[data_id].logit_bias, query_config.model_to_test
)
}
if prompts[data_id].logit_bias is not None
else {}
),
)
for data_id in prompts.keys()
]
model_responses = await asyncio.gather(*model_requests)
# pass through data that wasn't modified by the request
model_responses_map = {
data_id: response for data_id, response in zip(prompts.keys(), model_responses)
}
for key in data.keys():
if key not in prompts:
model_responses_map[key] = [data[key]]
response = format_response(data, model_responses_map, query_config)
return response
src/tools/__init__.py
View File
src/tools/dataloaders.py
-264
View File
@@ -1,264 +0,0 @@
__all__ = ["load_prompts", "load_problems", "load_problems_from_json", "load_solutions"]
import json
import logging
import os
from .path_utils import get_default_results_directory, get_root_directory
logger = logging.getLogger(__name__)
ROOT_DIR = get_root_directory()
DATA_DIR = ROOT_DIR / "data" / "APPS"
PROMPTS_DIR = ROOT_DIR / "src" / "prompts"
def get_data_dir():
return DATA_DIR
def load_prompts(prompt_type):
files = [f for f in (PROMPTS_DIR / prompt_type.value).glob("*") if f.is_file()]
prompts = {}
for file in files:
with file.open("r") as f:
prompts[file.name] = f.read()
return prompts
def load_problem_subset(subset, require_solutions=False, problem_ids=None):
def load_problems(dir, num_problems=None):
if problem_ids:
problem_dirs = problem_ids
else:
problem_dirs = os.listdir(dir)
problems = {}
added = 0
for problem_dir in problem_dirs:
problem_path = dir / problem_dir
problem = {}
with (problem_path / "metadata.json").open("r") as f:
problem["metadata"] = json.load(f)
if (
subset != "ALL"
and problem["metadata"]["difficulty"].lower() != subset.lower()
):
continue
if require_solutions and not (problem_path / "solutions.json").exists():
logger.debug(
f"Skipping problem {problem_dir} because it does not have solutions"
)
continue
with (problem_path / "question.txt").open("r") as f:
problem["question"] = f.read()
problem["uid"] = problem_dir
problems[problem_dir] = problem
added += 1
if added >= num_problems:
break
return problems
return load_problems
def load_problems(dir, num_problems=None):
return load_problem_subset("ALL")(dir, num_problems)
def load_problems_from_json(path, num_problems=None, problem_ids=None):
problems = {}
try:
with open(path) as f:
data = json.load(f)
except Exception as e:
print('read data error: ', e)
data = path
if num_problems is not None:
data = data[:num_problems]
for i, item in enumerate(data):
item["uid"] = i
if 'vanilla_label' not in item:
item["vanilla_label"] = item["label"]
problems[f"{i}"] = item
return problems
def load_problems_from_json_ids(path, num_problems=None, problem_ids=None):
problems = {}
try:
with open(path) as f:
data = json.load(f)
except:
data = path
if problem_ids is not None:
data = [data[i] for i in problem_ids]
for i, item in enumerate(data):
item["uid"] = i
if 'vanilla_label' not in item:
item["vanilla_label"] = item["label"]
# acc.append(item['label'] == item['vanilla_label'])
problems[f"{i}"] = item
return problems
def load_assignments(path, num_problems=None, problem_ids=None):
return path
def load_solutions(dir, num_problems=None):
solutions = {}
files = dir.glob("*")
files = [f for f in files if f.name != "incoming_problem_ids.json"]
if num_problems is not None:
files = list(files)[:num_problems]
for file in files:
with file.open("r") as f:
solution = json.load(f)
# Assume 1 solution per file
if isinstance(solution, list):
solutions[file.stem] = solution[0]
else:
solutions[file.stem] = solution
return solutions
def load_multiple_solutions(dir, num_problems=None, problem_ids=None):
solutions = {}
files = dir.glob("*")
files = [f for f in files if f.name != "incoming_problem_ids.json"]
for file in files:
with file.open("r") as f:
solution = json.load(f)
if "metadata" in solution:
solution.pop("metadata")
if "demonstration" in solution:
solution.pop("demonstration")
solutions[file.stem] = solution
return solutions
def load_multiple_solutions_w2s(dir, num_problems=None, problem_ids=None):
solutions = {}
files = dir.glob("*")
files = [f for f in files if f.name != "incoming_problem_ids.json"]
for file in files:
with file.open("r") as f:
solution = json.load(f)
metadata = solution[0]['metadata']
if "demonstration" in metadata:
metadata.pop("demonstration")
metadata['label'] = solution[0]['score'] > 0
solutions[file.stem] = metadata
return solutions
def save_to_cache(data, name, delete_existing=False, incoming_problem_ids=None):
dir = get_default_results_directory() / name
# Delete all files in the directory first
if delete_existing and os.path.exists(dir):
for file in os.listdir(dir):
file_path = os.path.join(dir, file)
if os.path.isfile(file_path):
os.unlink(file_path)
os.makedirs(dir, exist_ok=True)
for k, v in data.items():
if isinstance(v, list):
to_write = [
{
key: value
for key, value in item.items()
if key not in ["prompt", "response"]
}
for item in v
]
else:
to_write = {
key: value
for key, value in v.items()
if key not in ["prompt", "response"]
}
with open(dir / f"{k}.json", "w") as f:
json.dump(to_write, f, indent=4)
if incoming_problem_ids:
with open(dir / "incoming_problem_ids.json", "w") as f:
json.dump({"problem_ids": list(incoming_problem_ids)}, f, indent=4)
def read_from_cache(name):
dir = get_default_results_directory() / name
data = {}
incoming_problem_ids = []
for file in dir.glob("*.json"):
if file.name == "incoming_problem_ids.json":
with file.open("r") as f:
incoming_problem_ids = json.load(f).get("problem_ids", [])
else:
with file.open("r") as f:
value = json.load(f)
if not value.get("metadata"):
value["metadata"] = {k: v for k, v in value.items()}
data[file.stem] = value
return data, incoming_problem_ids
def load_ground_truth_solutions(problem_ids):
output = {}
for problem_id in problem_ids:
with open(get_data_dir() / "test" / problem_id / "solutions.json", "r") as f:
solutions = json.load(f)
cleaned_solutions = []
for solution in solutions:
# Remove unwanted lines from the solution
cleaned_solution = []
for line in solution.split("\n"):
if (
not line.strip().startswith("#!")
and " input=" not in line
and "sys.stdin" not in line
):
cleaned_solution.append(line)
cleaned_solutions.append("\n".join(cleaned_solution).strip())
output[problem_id] = cleaned_solutions
return output
def load_test_case(problem_id):
problem_id = problem_id.split("-")[0]
with open(get_data_dir() / "test" / problem_id / "input_output.json", "r") as f:
data = json.load(f)
return [
{"input": i, "output": o} for (i, o) in zip(data["inputs"], data["outputs"])
]
def load_test_cases(problem_ids):
output = {}
for problem_id in problem_ids:
output[problem_id] = load_test_case(problem_id)
return output
loaded_test_cases = {}
def get_test_cases_for_single_problem(problem_id):
global loaded_test_cases
if problem_id not in loaded_test_cases:
loaded_test_cases[problem_id] = load_test_case(problem_id)
return loaded_test_cases[problem_id]
src/tools/path_utils.py
-11
View File
@@ -1,11 +0,0 @@
__all__ = ["get_root_directory", "get_default_results_directory"]
from pathlib import Path
def get_root_directory():
return Path(__file__).parent.parent.parent
def get_default_results_directory():
return get_root_directory() / "results"
src/tools/printer.py
-134
View File
@@ -1,134 +0,0 @@
import json
import os
import subprocess
from pathlib import Path
from anytree import Node, RenderTree
from anytree.exporter import DotExporter
from src.tools.path_utils import get_default_results_directory
def print_experiment_log(experiment_name, strong_model, weak_model, problem_number):
results_dir = get_default_results_directory()
experiment_dir = results_dir / experiment_name
# Get all step directories
step_dirs = [d for d in experiment_dir.iterdir() if d.is_dir()]
# Sort step directories by step number and exclude "merged_results"
step_dirs = [d for d in step_dirs if d.name != "merged_results"]
step_dirs.sort(key=lambda x: int(x.name.split("-")[0]))
for step_dir in step_dirs:
# Check for both strong and weak model directories
for model in [strong_model, weak_model, f"{strong_model}+{weak_model}"]:
model_dir = step_dir / model
if not model_dir.exists():
continue
ignore_keys = ["metadata", "prompt", "response"]
if model == f"{strong_model}+{weak_model}":
ignore_keys.extend(["question", "test_cases", "uid"])
# Find matching problem files
problem_files = list(model_dir.glob(f"{problem_number}*.json"))
for problem_file in problem_files:
with open(problem_file, "r") as f:
data = json.load(f)
print(f"Step: {step_dir.name}")
print(f"Model: {model}")
print(f"Problem: {problem_file.stem}")
if not isinstance(data, list):
data = [data]
print("\nPrompt:")
prompt_array = data[0].get("prompt")
if prompt_array is None:
print("No prompt available")
else:
for text in prompt_array:
print(f"Role: {text['role']}")
print(f"Content: {text['content']}")
for response in data:
print("\nResponse:")
print(
response.get("response", {}).get(
"completion", "No response available"
)
)
print("\nOther Fields:")
for key, value in response.items():
if key not in ignore_keys:
print(f"{key}: {value}")
print("\n" + "=" * 50 + "\n")
def show_pipeline_graph(pipeline):
import matplotlib.pyplot as plt
import networkx as nx
# Create a directed graph
G = nx.DiGraph()
# Add nodes and edges
for task in pipeline.steps:
G.add_node(task.name)
for dep in task.dependencies:
G.add_edge(dep.name, task.name)
# Print the graph structure
print("Pipeline Dependency Graph:")
for node in nx.topological_sort(G):
predecessors = list(G.predecessors(node))
successors = list(G.successors(node))
print(f"{node}:")
if predecessors:
print(f" Parents: {', '.join(predecessors)}")
if successors:
print(f" Children: {', '.join(successors)}")
# Generate a DOT file for visualization
output_dir = get_default_results_directory() / pipeline.config.name
output_dir.mkdir(parents=True, exist_ok=True)
dot_file = output_dir / "pipeline_graph.dot"
png_file = output_dir / "pipeline_graph.png"
nx.drawing.nx_pydot.write_dot(G, str(dot_file))
print(f"DOT file generated at: {dot_file}")
# Generate PNG file using Graphviz
try:
subprocess.run(["dot", "-Tpng", str(dot_file), "-o", str(png_file)], check=True)
print(f"PNG file generated at: {png_file}")
except subprocess.CalledProcessError:
print(
"Error: Failed to generate PNG. Make sure Graphviz is installed and accessible in your PATH."
)
except FileNotFoundError:
print(
"Error: Graphviz not found. Please install Graphviz to generate PNG files."
)
# Optionally, you can also use matplotlib to visualize the graph
plt.figure(figsize=(12, 8))
pos = nx.spring_layout(G)
nx.draw(
G,
pos,
with_labels=True,
node_color="lightblue",
node_size=2000,
font_size=8,
arrows=True,
)
plt.title("Pipeline Dependency Graph")
plt.axis("off")
plt.tight_layout()
plt.savefig(str(output_dir / "pipeline_graph_matplotlib.png"))
print(
f"Matplotlib graph generated at: {output_dir / 'pipeline_graph_matplotlib.png'}"
)
src/tools/string_manipulation.py
-62
View File
@@ -1,62 +0,0 @@
import re
def format_key_suffix(key_suffix):
if key_suffix:
if key_suffix[0] != "_":
key_suffix = f"_{key_suffix}"
else:
key_suffix = ""
return key_suffix
COMMENTS_REGEX = r"""
^ # Begin of line.
(?:
# A) Capturing group n°1: Full-line comment followed by empty lines.
(
[ \t]* # Optional spaces or tabs.
\#[^\r\n]*\r?\n # The comment and the new line.
(?:[ \t]*\r?\n)* # Optional empty lines (perhaps with spaces/tabs).
)
|
# B) Statement and optional comment at the end.
(?:
( # Capturing group n°2 : The statement
(?:
# Multi-line strings with \"\"\" or '''.
# Capturing group n°3 : The triple quotes.
(['\"]{3})[\s\S]*?\3
|
# Double-quoted string "It's ok".
\"(?: \\. | [^\"] )*\"
|
# Single-quoted string 'I\'ll say "Hello!"'.
'(?: \\. | [^'] )*'
|
# Any chars, except spaces, hashtag, quotes and new lines.
[^ \t#\"'\r\n]+
|
# Horizontal spaces, but not followed by a comment, because
# we want the spaces in front of the comment to be matched
# together with the optional comment we want to get rid of.
[ \t]+(?![ \t]*\#)
)+
)
# Capturing group n°4: An optional comment at the end of a statement.
(
[ \t]*\#[^\r\n]*
)?
)+
)
"""
def strip_comments_from_string(string):
return re.sub(
COMMENTS_REGEX,
"\\2",
string,
0,
re.MULTILINE | re.VERBOSE | re.UNICODE,
)
src/utils.py
-239
View File
@@ -1,239 +0,0 @@
import asyncio
import json
import logging
import os
import time
from functools import lru_cache, wraps
from pathlib import Path
from typing import Callable
import matplotlib.pyplot as plt
import numpy as np
import openai
import pandas as pd
import replicate
import typer
import yaml
from tenacity import retry, retry_if_result, stop_after_attempt
typer.main.get_command_name = lambda name: name
LOGGER = logging.getLogger(__name__)
SEPARATOR = "---------------------------------------------\n\n"
SEPARATOR_CONVERSATIONAL_TURNS = "=============================================\n\n"
PROMPT_HISTORY = "prompt_history"
SECRETS_FILE_PATH = Path(__file__).parent.parent / "SECRETS"
LOGGING_LEVELS = {
"critical": logging.CRITICAL,
"error": logging.ERROR,
"warning": logging.WARNING,
"info": logging.INFO,
"debug": logging.DEBUG,
}
def setup_environment(
anthropic_tag: str = "ANTHROPIC_API_KEY",
logger_level: str = "info",
openai_tag: str = "API_KEY",
mistral_tag: str = "MISTRAL_API_KEY",
replicate_tag: str = "REPLICATE_API_KEY",
organization: str = None,
):
setup_logging(logger_level)
load_secrets(
SECRETS_FILE_PATH,
anthropic_tag,
logger_level,
openai_tag,
mistral_tag,
replicate_tag,
organization,
)
def setup_logging(level_str):
level = LOGGING_LEVELS.get(
level_str.lower(), logging.INFO
) # default to INFO if level_str is not found
logging.basicConfig(
level=level,
format="%(asctime)s [%(levelname)s] (%(name)s) %(message)s",
datefmt="%Y-%m-%d %H:%M:%S",
)
root_logger = logging.getLogger()
root_logger.setLevel(level)
# Disable logging from noisy libraries
logging.getLogger("openai").setLevel(logging.CRITICAL)
logging.getLogger("httpx").setLevel(logging.CRITICAL)
logging.getLogger("matplotlib").setLevel(logging.CRITICAL)
logging.getLogger("anthropic").setLevel(logging.CRITICAL)
logging.getLogger("httpcore").setLevel(logging.CRITICAL)
logging.getLogger("urllib3").setLevel(logging.CRITICAL)
LOGGER.info(f"Logging level set to {level_str}")
def load_secrets(
file_path=SECRETS_FILE_PATH,
anthropic_tag: str = "ANTHROPIC_API_KEY",
logger_level: str = "info",
openai_tag: str = "API_KEY",
mistral_tag: str = "MISTRAL_API_KEY",
replicate_tag: str = "REPLICATE_API_KEY",
organization: str = None,
):
secrets = {}
with open(file_path) as f:
for line in f:
key, value = line.strip().split("=", 1)
secrets[key] = value
openai.api_key = secrets[openai_tag]
os.environ['LLAMA_API_BASE'] = secrets['LLAMA_API_BASE']
# replicate.api_token = secrets[replicate_tag]
# os.environ["ANTHROPIC_API_KEY"] = secrets[anthropic_tag]
# os.environ["MISTRAL_API_KEY"] = secrets[mistral_tag]
# os.environ["REPLICATE_API_KEY"] = secrets[replicate_tag]
if organization is not None:
openai.organization = secrets[organization]
if secrets.get("API_BASE") is not None:
openai.api_base = secrets['API_BASE']
return secrets
def load_yaml(file_path):
with open(file_path) as f:
content = yaml.safe_load(f)
return content
@lru_cache(maxsize=8)
def load_yaml_cached(file_path):
with open(file_path) as f:
content = yaml.safe_load(f)
return content
def save_yaml(file_path, data):
with open(file_path, "w") as f:
yaml.dump(data, f)
def load_jsonl(file_path):
data = []
with open(file_path, "r") as f:
for line in f:
json_obj = json.loads(line)
data.append(json_obj)
return data
def save_jsonl(file_path, data):
with open(file_path, "w") as f:
for line in data:
json.dump(line, f)
f.write("\n")
def delete_old_prompt_files(
path: str = PROMPT_HISTORY, max_age_minutes: int = 60, keep_recent: int = 50
):
"""
Delete all files in the folder that:
- Are more than max_age_minutes old
- AND are not one of the keep_recent most recent files
"""
if not os.path.exists(path):
return
# Get all files in the folder with their full paths and creation times
files = [
{
"path": os.path.join(path, filename),
"ctime": os.path.getctime(os.path.join(path, filename)),
}
for filename in os.listdir(path)
if os.path.isfile(os.path.join(path, filename))
]
# Sort files by creation time
files.sort(key=lambda f: f["ctime"], reverse=True)
# Current time in seconds since epoch
now = time.time()
deleted_count = 0
for index, file_info in enumerate(files):
# File age in minutes
age_minutes = (now - file_info["ctime"]) / 60
# If file is older than x_minutes and is not one of the y_most_recent files, delete it
if age_minutes > max_age_minutes and index >= keep_recent:
os.remove(file_info["path"])
deleted_count += 1
if deleted_count > 0:
print(f"Deleted {deleted_count} old prompt files")
def typer_async(f):
@wraps(f)
def wrapper(*args, **kwargs):
try:
loop = asyncio.get_running_loop()
except RuntimeError: # No event loop running
loop = None
if loop is None:
return asyncio.run(f(*args, **kwargs))
else:
return f(*args, **kwargs) # Return coroutine to be awaited
return wrapper
@retry(
stop=stop_after_attempt(16),
retry=retry_if_result(lambda result: result is not True),
)
def function_with_retry(function, *args, **kwargs):
return function(*args, **kwargs)
@retry(
stop=stop_after_attempt(16),
retry=retry_if_result(lambda result: result is not True),
)
async def async_function_with_retry(function, *args, **kwargs):
return await function(*args, **kwargs)
def log_model_timings(api_handler, save_location="./model_timings.png"):
if len(api_handler.model_timings) > 0:
plt.figure(figsize=(10, 6))
for model in api_handler.model_timings:
timings = np.array(api_handler.model_timings[model])
wait_times = np.array(api_handler.model_wait_times[model])
LOGGER.info(
f"{model}: response {timings.mean():.3f}, waiting {wait_times.mean():.3f} (max {wait_times.max():.3f}, min {wait_times.min():.3f})"
)
plt.plot(
timings, label=f"{model} - Response Time", linestyle="-", linewidth=2
)
plt.plot(
wait_times, label=f"{model} - Waiting Time", linestyle="--", linewidth=2
)
plt.legend()
plt.title("Model Performance: Response and Waiting Times")
plt.xlabel("Sample Number")
plt.ylabel("Time (seconds)")
plt.savefig(save_location, dpi=300)
plt.close()
def softmax(x):
return np.exp(x) / np.sum(np.exp(x), axis=0)