diff --git a/docs/icm_progress.png b/docs/icm_progress.png deleted file mode 100644 index b3dcf38..0000000 Binary files a/docs/icm_progress.png and /dev/null differ diff --git a/icm_progress.png b/icm_progress.png index 0742693..2b54e1d 100644 Binary files a/icm_progress.png and b/icm_progress.png differ diff --git a/pyproject.toml b/pyproject.toml index 8a463c3..5462645 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -33,6 +33,7 @@ dependencies = [ "scikit-learn>=1.7.2", "scipy>=1.15.3", "seaborn>=0.13.2", + "simple-parsing>=0.1.7", "sqlalchemy==2.0.18", "tabulate>=0.9.0", "tenacity>=9.1.2", diff --git a/src/data/daily_dilemma.py b/src/data/daily_dilemma.py deleted file mode 100644 index 0e329c7..0000000 --- a/src/data/daily_dilemma.py +++ /dev/null @@ -1,24 +0,0 @@ -from datasets import load_dataset - -def load_daily_dilemma(): - dataset = load_dataset("kellycyy/daily_dilemmas", split="test") - - data = [] - for item in dataset: - # Unbiased prompt - just the situation and action - prompt = f"Situation: {item['dilemma_situation']}\nAction: {item['action']}" - - # Map to A/B (unrelated group names) - label = 1 if item['action_type'] == 'to_do' else 0 - consistency_key = 'A' if item['action_type'] == 'to_do' else 'B' - - example = { - 'uid': item['idx'], - 'prompt': prompt, - 'vanilla_label': label, - 'consistency_id': item['dilemma_idx'], - 'consistency_key': consistency_key - } - data.append(example) - - return data diff --git a/src/data/daily_dilemmas.py b/src/data/daily_dilemmas.py new file mode 100644 index 0000000..c5fb707 --- /dev/null +++ b/src/data/daily_dilemmas.py @@ -0,0 +1,96 @@ +from datasets import load_dataset +import pandas as pd +from collections import defaultdict +import ast + + +def convert_values_to_list(x): + # turn into list + s = x["values_aggregated"] + v = ast.literal_eval(s) + return {"values_aggregated": v} + +def load_daily_dilemma(label_col = "label_Virtue/Truthfulness"): + dataset = load_dataset("kellycyy/daily_dilemmas", split="test") + dataset = dataset.map(convert_values_to_list) + df_labels = load_labels(dataset).rename(columns=lambda x: f"label_{x}" if x != "dilemma_idx" else x) + df = dataset.to_pandas() + df = df.merge(df_labels, on='dilemma_idx') + cols_label = [c for c in df_labels.columns if c.startswith('label_')] + # flip labels on not_to_do + df.loc[df['action_type'] == 'not_to_do', cols_label] *= -1 + dataset = df.to_dict(orient='records') + + + + data = [] + for item in dataset: + # Unbiased prompt - just the situation and action + prompt = f"Situation: {item['dilemma_situation']}\nAction: {item['action']}" + + # Map to A/B (unrelated group names) + consistency_key = 'A' if item['action_type'] == 'to_do' else 'B' + + example = { + 'uid': item['idx'], + 'prompt': prompt, + 'vanilla_label': item[label_col], + 'consistency_id': item['dilemma_idx'], + 'consistency_key': consistency_key + } + data.append(example) + + + return data + + + +def load_labels(dd_dataset): + ds_values = load_dataset("kellycyy/daily_dilemmas", split="test", name="Values") + + # moral tags + moral_frameworks = ["WVS", "MFT", "Virtue", "Emotion", "Maslow"] + + value2framework_dicts = {} + for framework in moral_frameworks: + df_values = ds_values.to_pandas()[["value", framework]].dropna() + value2framework_dict = df_values.set_index("value")[framework].to_dict() + value2framework_dict = {k: f"{framework}/{v}" for k, v in value2framework_dict.items()} + value2framework_dicts[framework] = value2framework_dict + + + # make labels + df_dilemma = dd_dataset.to_pandas()[["dilemma_idx", "action_type", "values_aggregated"]] + dilemma_idx = df_dilemma["dilemma_idx"].unique() + + labels = [] + for d_idx in dilemma_idx: + pos_values = ( + df_dilemma.query('dilemma_idx == @d_idx and action_type == "to_do"')["values_aggregated"].iloc[0].tolist() + ) + neg_values = ( + df_dilemma.query('dilemma_idx == @d_idx and action_type == "not_to_do"')["values_aggregated"].iloc[0].tolist() + ) + + label = defaultdict(int) + + for framework in value2framework_dicts: + value2framework_dict = value2framework_dicts[framework] + virtues = sorted(set(value2framework_dict.values())) + + pos_virtues = [value2framework_dict[k] for k in pos_values if k in value2framework_dict] + neg_virtues = [value2framework_dict[k] for k in neg_values if k in value2framework_dict] + + for p in pos_virtues: + label[p] += 1 + for n in neg_virtues: + label[n] -= 1 + + labels.append(dict(dilemma_idx=d_idx, **label)) + + + + df_labels = pd.DataFrame(labels).set_index("dilemma_idx") + assert df_labels.index.is_unique + return df_labels + diff --git a/src/simple_icm.py b/src/simple_icm.py index 181fc3e..913d047 100644 --- a/src/simple_icm.py +++ b/src/simple_icm.py @@ -19,7 +19,7 @@ from dataclasses import dataclass, asdict import dotenv from loguru import logger from openrouter_wrapper.logprobs import openrouter_completion_wlogprobs, get_logprobs_choices, LogprobsNotSupportedError # User's wrapper -from typing import List, Tuple, Callable +from typing import List, Tuple, Callable, Literal import asyncio from aiocache import cached from itertools import combinations @@ -43,6 +43,7 @@ logger.add(sys.stderr, format="{time:YYYY-MM-DD HH:mm} | { # Global cost tracker total_cost = 0.0 +reasoning_log = "" # %% [code] @@ -56,16 +57,22 @@ class Config: num_seed: int = 8 max_iters: int = 2500 # should be at least dataset size X 2 n_shots: int = 6 # Number of in-context examples - model_id: str = "meta-llama/llama-3.1-8b-instruct" # Logprobs supported - provider_whitelist: Tuple[str] = None # None to let OpenRouter choose + # model_id: str = "meta-llama/llama-3.1-8b-instruct" # Logprobs supported + model_id: str = "qwen/qwen3-235b-a22b-2507" # Logprobs supported + provider_whitelist: Tuple[str] = ('Chutes','Nebius',) # None to let OpenRouter choose out_dir: Path = Path("../outputs/icm") # Directory to save outputs log_interval: int = 100 # Log progress every N iterations + dataset: Literal["truthfulqa", "daily_dilemmas"] = "truthfulqa" # Dataset name for logging -C = Config( - model_id="qwen/qwen3-235b-a22b-2507", # $0.2 0.6 - provider_whitelist=[ 'Chutes','Nebius',], -) -C.out_dir.mkdir(parents=True, exist_ok=True) + +import simple_parsing + +C: Config = simple_parsing.parse(Config) + +# C = Config( +# model_id="qwen/qwen3-235b-a22b-2507", # $0.2 0.6 +# provider_whitelist=[ 'Chutes','Nebius',], +# ) # C = Config( # model_id="qwen/qwen3-30b-a3b-instruct-2507", # 0.08 $0.33 @@ -79,14 +86,23 @@ C.out_dir.mkdir(parents=True, exist_ok=True) logger.info(f"Config: {C}") config_dict = asdict(C) -config_dict['out_dir'] = str(config_dict['out_dir']) -with open(C.out_dir / "icm_config.json", "w") as f: +out_dir = C.out_dir / C.dataset.replace(' ', '_').replace('/', '_') +out_dir.mkdir(parents=True, exist_ok=True) +config_dict['out_dir'] = str(C.out_dir) +with open(out_dir / "icm_config.json", "w") as f: json.dump(config_dict, f, indent=2) # %% [code] from src.data.truthfulqa import load_truthfulqa, is_consistent +from data.daily_dilemmas import load_daily_dilemma + +if C.dataset == "truthfulqa": + data = load_truthfulqa() +elif C.dataset == "daily_dilemmas": + data = load_daily_dilemma() +else: + raise ValueError(f"Unknown dataset {C.dataset}") -data = load_truthfulqa() logger.info("Loaded {} examples", len(data)) # %% [code] @@ -145,6 +161,9 @@ async def predict_label(example_uid, current_demos, config=C, verbose=False, all {"role": "user", "content": instruction+"".join(fewshot)+f"\n\n## Candidate:\n{target_prompt}"}, {"role": "assistant", "content": "\n## Set:"} # Assistant prefill to ensure ] + + if verbose>1: + messages[0]['content'] = "ALWAYS GIVE BRIEF REASONING AFTERWARDS. " + messages[0]['content'] response = await cached_openrouter_completion_wlogprobs( @@ -152,6 +171,7 @@ async def predict_label(example_uid, current_demos, config=C, verbose=False, all provider_whitelist=config.provider_whitelist, messages=messages, max_completion_tokens=160 if verbose else 5, + min_completion_tokens=30 if verbose else 1, temperature=0.4, top_logprobs=8, ) @@ -160,9 +180,14 @@ async def predict_label(example_uid, current_demos, config=C, verbose=False, all if verbose: logger.info(f"Debug Prediction - UID {example_uid}:") - logger.info(f"messages: {print_messages(messages)}") - logger.info(f"Response Content: {response['choices'][0]['message']['content']}") + logger.info(f"messages: `{print_messages(messages)}`") + logger.info(f"Response Content: `{response['choices'][0]['message']['content']}`") logger.info(f"--- End Debug ---") + if verbose>1: + global reasoning_log + reasoning_log += f"\n\n## Candidate:\n{target_prompt}\n## Set:\n" + reasoning_log += response['choices'][0]['message']['content'] + try: choice_strs = ["A", "B"] @@ -174,7 +199,7 @@ async def predict_label(example_uid, current_demos, config=C, verbose=False, all if not choice_in_toplogp: model_response = response['choices'][0]['message']['content'] - logger.warning(f"Choices not returned for UID {example_uid}, may indicate model confusion. choice_logp={choice_logp}. Instead we got these top logprobs: {top_logp} and \nmessages: {print_messages(messages)}\nthis output: {model_response}") + logger.warning(f"Choices not returned for UID {example_uid}, may indicate model confusion. choice_logp={choice_logp}. Instead we got these top logprobs: {top_logp} and \nmessages: ...`{print_messages(messages)[-90:]}`\nthis output:`{model_response}`") score = choice_logp["A"] - choice_logp["B"] predicted = 1 if score > 0 else 0 return predicted, float(score) @@ -223,6 +248,9 @@ def compute_energy(demos, config=C): energy = config.alpha * avg_lprob - num_inconsistent - (num_inconsistent / max(1, len(labeled))) # Normalized penalty accuracy = np.mean([d['label'] == d['vanilla_label'] for d in labeled]) + # flip acc if needed, as this is unsupervised + if accuracy < 0.5: + accuracy = 1 - accuracy return energy, { 'avg_lprob': avg_lprob, 'num_inconsistent': num_inconsistent, @@ -405,6 +433,8 @@ async def run_icm(demonstrations, config=C): except KeyboardInterrupt: logger.info("Stopping early.") + except asyncio.CancelledError: + logger.info("Asyncio task cancelled.") return demonstrations, energies, accuracies @@ -418,13 +448,13 @@ logger.info("\nFinal Results:") logger.info("Total cost: ${:.4f}", total_cost) logger.info("Energy: {:.2f}", final_energy) # TODO show vanilla accuracy here for comparison -logger.info("Accuracy vs vanilla: {:.2f}, initial {:.2f}", final_metrics['accuracy'], accuracies[0]) +logger.info("Accuracy [labelled] vs vanilla: {:.2f}, initial {:.2f}", final_metrics['accuracy'], accuracies[0]) logger.info("Labeled: {}/{}", final_metrics['num_labeled'], len(data)) logger.info("Inconsistencies: {}", final_metrics['num_inconsistent']) # Final labels df = pd.DataFrame(final_demos).T -df.to_parquet(C.out_dir / "icm_final_labels.parquet") +df.to_parquet(out_dir / "icm_final_labels.parquet") df_labeled = df.dropna(subset='label').sort_values(by='score', key=np.abs, ascending=False) df_labeled_disagreed = df_labeled[df_labeled['vanilla_label'] != df_labeled['label']] @@ -435,11 +465,14 @@ print(df_labeled_disagreed[['consistency_id', 'label', 'vanilla_label', 'score', for uid, row in df_labeled_disagreed.iterrows(): print(f"\n## Candidate: {row['prompt']}\nICM Set: {'A' if row['label']==1 else 'B'}, Vanilla Set: {'A' if row['vanilla_label']==1 else 'B'}, score={row['score']}\n") - +print(f"\nFinal labeled examples saved to {out_dir / 'icm_final_labels.parquet'}") # %% [code] # Simple visualization (requires matplotlib) +with open(out_dir / "cost.txt", "w") as f: + f.write(f"\n\nTotal cost: ${total_cost:.4f}\n") + f.write(reasoning_log) plt.figure(figsize=(10, 4)) plt.subplot(1, 2, 1) @@ -450,7 +483,7 @@ plt.ylabel('Energy') plt.subplot(1, 2, 2) plt.plot(accuracies) -plt.title('Accuracy over Iterations') +plt.title('Accuracy [labelled] over Iterations') plt.xlabel('Iteration') plt.ylabel('Accuracy') diff --git a/uv.lock b/uv.lock index 1bae939..210e8bc 100644 --- a/uv.lock +++ b/uv.lock @@ -3452,6 +3452,19 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/e0/f9/0595336914c5619e5f28a1fb793285925a8cd4b432c9da0a987836c7f822/shellingham-1.5.4-py2.py3-none-any.whl", hash = "sha256:7ecfff8f2fd72616f7481040475a65b2bf8af90a56c89140852d1120324e8686", size = 9755, upload-time = "2023-10-24T04:13:38.866Z" }, ] +[[package]] +name = "simple-parsing" +version = "0.1.7" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "docstring-parser" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/eb/c5/f1e2fcb3a81085cdf3cfed48b8c8ce0e7cc30c95dee734cbb35d6265336a/simple_parsing-0.1.7.tar.gz", hash = "sha256:225e6b35252d68f7894716101fe3bd7e6dd3d30ab7b1c3c023f77a42dbe1336f", size = 96375, upload-time = "2025-01-20T19:46:35.986Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/4f/9c/e9ea38750027a6de3e3c5e68a19fda0e7b0cd3db8045f30d0f6bc113b911/simple_parsing-0.1.7-py3-none-any.whl", hash = "sha256:5276e6c90c157362dd0173d1eecebe58361a66b457129cc9bba13b78a4e85092", size = 112782, upload-time = "2025-01-20T19:46:33.325Z" }, +] + [[package]] name = "six" version = "1.17.0" @@ -3782,6 +3795,7 @@ dependencies = [ { name = "scipy", version = "1.15.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, { name = "scipy", version = "1.16.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, { name = "seaborn" }, + { name = "simple-parsing" }, { name = "sqlalchemy" }, { name = "tabulate" }, { name = "tenacity" }, @@ -3827,6 +3841,7 @@ requires-dist = [ { name = "scikit-learn", specifier = ">=1.7.2" }, { name = "scipy", specifier = ">=1.15.3" }, { name = "seaborn", specifier = ">=0.13.2" }, + { name = "simple-parsing", specifier = ">=0.1.7" }, { name = "sqlalchemy", specifier = "==2.0.18" }, { name = "tabulate", specifier = ">=0.9.0" }, { name = "tenacity", specifier = ">=9.1.2" },