diff --git a/nbs/simple_icm.py b/nbs/simple_icm.py index bddcc2f..ae680e3 100644 --- a/nbs/simple_icm.py +++ b/nbs/simple_icm.py @@ -69,21 +69,36 @@ choice_logprobs_permuted_d # %% [code] -# Small hardcoded dataset: Math claims (inspired by GSM8K) -# Each group shares 'consistency_id' (question), 'consistency_key' (claim answer) -data = [ - # Group 1: Question "What is 2+2?" - {'uid': 0, 'prompt': 'Question: What is 2+2? Claim: The answer is 4. I think this claim is ', 'vanilla_label': 1, 'consistency_id': 'q1', 'consistency_key': '4'}, - {'uid': 1, 'prompt': 'Question: What is 2+2? Claim: The answer is 5. I think this claim is ', 'vanilla_label': 0, 'consistency_id': 'q1', 'consistency_key': '5'}, - {'uid': 2, 'prompt': 'Question: What is 2+2? Claim: The answer is 3. I think this claim is ', 'vanilla_label': 0, 'consistency_id': 'q1', 'consistency_key': '3'}, - {'uid': 3, 'prompt': 'Question: What is 2+2? Claim: The answer is 4. Final check: True. I think this claim is ', 'vanilla_label': 1, 'consistency_id': 'q1', 'consistency_key': '4'}, +from datasets import load_dataset + +# Load larger HF dataset: Yik/truthfulQA-bool +dataset = load_dataset("Yik/truthfulQA-bool", split="train") +data = [] +group_id = 0 +groups = {} +for idx, item in enumerate(dataset): + claim = item['question'] # The question itself is the claim + label = item['label'] + consistency_key = str(claim)[:10] # Short unique key + consistency_id = group_id - # Group 2: Question "What is 3*3?" - {'uid': 4, 'prompt': 'Question: What is 3*3? Claim: The answer is 9. I think this claim is ', 'vanilla_label': 1, 'consistency_id': 'q2', 'consistency_key': '9'}, - {'uid': 5, 'prompt': 'Question: What is 3*3? Claim: The answer is 6. I think this claim is ', 'vanilla_label': 0, 'consistency_id': 'q2', 'consistency_key': '6'}, - {'uid': 6, 'prompt': 'Question: What is 3*3? Claim: The answer is 10. I think this claim is ', 'vanilla_label': 0, 'consistency_id': 'q2', 'consistency_key': '10'}, - {'uid': 7, 'prompt': 'Question: What is 3*3? Claim: The answer is 9. Final check: True. I think this claim is ', 'vanilla_label': 1, 'consistency_id': 'q2', 'consistency_key': '9'}, -] + # Simple grouping: Create fake variations per group (e.g., 4 claims per 80 groups for batch) + if idx % 4 == 0: + group_id += 1 + groups[group_id] = [] + + example = { + 'uid': idx, + 'prompt': f"Question: {claim}\nClaim: {claim}\nI think this claim is ", + 'vanilla_label': label, + 'consistency_id': consistency_id, + 'consistency_key': consistency_key, + } + data.append(example) + +# Limit to small batch for demo (e.g., 64 items) +data = data[:64] +logger.info(f"Loaded {len(data)} examples from TruthfulQA-bool") # %% [code] # Initialize: Random labels for first num_seed, None for others @@ -104,16 +119,23 @@ logger.info("Initialized labels: {}", {k: v['label'] for k, v in demonstrations. # Predict label using in-context prompting (placeholder with OpenAI) def predict_label(example_uid, current_demos, config=C): - # Simplified few-shot prompt (distilled from original get_judge_prompt_fewshot) - # TODO: Further refine by inspecting src/model_querying/prompt_creation.py for better template (e.g., full task description, more examples) + # Refined few-shot prompt based on get_judge_prompt_fewshot + # Sort by consistency_key for relevance, limit to 8 + relevant_demos = sorted( + [demo for uid, demo in current_demos.items() if uid != example_uid and demo['label'] is not None], + key=lambda x: x.get('consistency_key', '') + )[:8] + + instruction = """You are a judge evaluating the truthfulness of claims based on factual knowledge. Look at the examples below and judge if each claim is True (1) or False (0). + +""" fewshot = [] - for uid, demo in current_demos.items(): - if uid != example_uid and demo['label'] is not None: - label_str = "1" if demo['label'] == 1 else "0" # Binary for logprobs - fewshot.append(f"{demo['prompt']} {label_str}.") + for idx, demo in enumerate(relevant_demos): + label_str = "1" if demo['label'] == 1 else "0" + fewshot.append(f"### Example {idx + 1}:\n{demo['prompt']}Judgment: {label_str}\n\n") target_prompt = demonstrations[example_uid]['prompt'] - full_prompt = "Classify claims as true (1) or false (0) based on correctness. Examples:\n" + "\n".join(fewshot) + f"\n{target_prompt}" + full_prompt = instruction + "".join(fewshot) + target_prompt + "Judgment:" try: # Use wrapper for chat completion (messages format) @@ -232,13 +254,41 @@ def run_icm(demonstrations, config=C): old_energy, old_metrics = compute_energy(demonstrations, config) for iter in range(config.max_iters): - # Sample example (simple random; TODO: weight by group inconsistencies) + # Weighted sampling for inconsistent groups + groups = {} all_uids = list(demonstrations.keys()) - example_uid = random.choice(all_uids) + for uid in all_uids: + cid = demonstrations[uid]['consistency_id'] + if cid not in groups: + groups[cid] = [] + groups[cid].append(uid) - if example_uid in current_labeled and demonstrations[example_uid]['label'] is not None: - # Can re-label existing - pass + weights = [0.1 for _ in all_uids] # Base low + + for cid, group_uids in groups.items(): + labeled_labels = [demonstrations[uid]['label'] for uid in group_uids if demonstrations[uid]['label'] is not None] + num_labeled = len(labeled_labels) + num_unlabeled = len(group_uids) - num_labeled + + inconsistency = 0 + if num_labeled > 0: + unique_labels = set(labeled_labels) + inconsistency = 1 if len(unique_labels) > 1 else 0 + + if num_unlabeled > 0: + weight_factor = (0.5 + 0.5 * inconsistency) * (1 + num_unlabeled / len(group_uids)) + for uid in group_uids: + if demonstrations[uid]['label'] is None: # Unlabeled + idx = all_uids.index(uid) + weights[idx] = weight_factor + else: + # Fully labeled low priority + for uid in group_uids: + idx = all_uids.index(uid) + weights[idx] = 0.1 + + weights = [max(w, 0.01) for w in weights] + example_uid = random.choices(all_uids, weights=weights)[0] # Predict new label new_label, score = predict_label(example_uid, current_labeled, config) @@ -312,9 +362,10 @@ plt.show() # %% [markdown] # ## Next Steps +# - [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: Add weighted sampling for inconsistent groups (currently random). -# - FIXME: Enhance consistency fix to multi-iter proposals like ICM.py. -# - TODO: Load larger HuggingFace dataset (e.g., 'Yik/truthfulQA-bool' subset, format to messages). -# - TODO: Refine few-shot prompt from original get_judge_prompt_fewshot. -# - Test with logprobs-supported models. +# - FIXME: +# th logprobs-supported models. + diff --git a/pyproject.toml b/pyproject.toml index ed11fd6..81cba5e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -18,6 +18,7 @@ dependencies = [ "fastapi==0.100.0", "fire>=0.7.1", "hydra-core>=1.3.2", + "loguru>=0.7.3", "matplotlib>=3.10.6", "mistralai>=1.9.10", "openai==0.28.0", diff --git a/uv.lock b/uv.lock index 9901d26..55d6de3 100644 --- a/uv.lock +++ b/uv.lock @@ -3751,6 +3751,7 @@ dependencies = [ { name = "fastapi" }, { name = "fire" }, { name = "hydra-core" }, + { name = "loguru" }, { name = "matplotlib" }, { name = "mistralai" }, { name = "openai" }, @@ -3795,6 +3796,7 @@ requires-dist = [ { name = "fastapi", specifier = "==0.100.0" }, { name = "fire", specifier = ">=0.7.1" }, { name = "hydra-core", specifier = ">=1.3.2" }, + { name = "loguru", specifier = ">=0.7.3" }, { name = "matplotlib", specifier = ">=3.10.6" }, { name = "mistralai", specifier = ">=1.9.10" }, { name = "openai", specifier = "==0.28.0" },