work with truthfulQA

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.
+

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",

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" },