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

vibe

Browse Source
This commit is contained in:
wassname
2025-10-04 15:18:39 +08:00
parent 7cd4bf48dc
commit 2132333b4d
3 changed files with 91 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files
nbs/simple_icm.py
+79 -10
View File
@@ -23,6 +23,18 @@ from loguru import logger
from openrouter_wrapper.logprobs import openrouter_completion_wlogprobs, get_logprobs_choices, LogprobsNotSupportedError # User's wrapper
from typing import List, Tuple
import asyncio
from aiocache import SimpleMemoryCache, cached, Cache
cache = SimpleMemoryCache()
try:
from IPython import get_ipython
if get_ipython() is not None: # In Jupyter/VS Code notebook
import nest_asyncio
nest_asyncio.apply()
except ImportError:
pass # Not in notebook or IPython not available
dotenv.load_dotenv()
@@ -43,13 +55,13 @@ class Config:
num_seed: int = 8
max_iters: int = 950 # Small for demo; increase for more
n_shots: int = 16 # Number of in-context examples
model: str = "meta-llama/llama-3.1-8b-instruct" # Logprobs supported
model_id: str = "meta-llama/llama-3.1-8b-instruct" # Logprobs supported
provider_whitelist: Tuple[str] = None # None to let OpenRouter choose
out_dir: Path = Path("../outputs/icm") # Directory to save outputs
log_interval: int = 50 # Log progress every N iterations
C = Config(
model="qwen/qwen3-235b-a22b-2507",
model_id="qwen/qwen3-235b-a22b-2507",
provider_whitelist=[ 'Chutes','Nebius',],
)
C.out_dir.mkdir(parents=True, exist_ok=True)
@@ -120,15 +132,16 @@ logger.info("Initialized labels: {}", {k: v['label'] for k, v in demonstrations.
def print_messages(messages):
return "\n".join([f"**{m['role'].upper()}**: {m['content']}" for m in messages])
@cached(cache)
async def cached_openrouter_completion_wlogprobs(*args, **kwargs):
# TODO: Async caching (e.g., with aiocache)
return await openrouter_completion_wlogprobs(*args, **kwargs)
async def predict_label(example_uid, current_demos, config=C, verbose=False):
async def predict_label(example_uid, full_demos, labeled_demos, config=C, verbose=False):
# simplified from https://github.com/Jiaxin-Wen/Unsupervised-Elicitation/blob/master/src/experiments/ICM.py#L254
# Sort by consistency_key for relevance, limit to 16
relevant_demos = sorted(
[demo for uid, demo in current_demos.items() if uid != example_uid and demo['label'] is not None],
[demo for uid, demo in labeled_demos.items() if uid != example_uid],
key=lambda x: x.get('consistency_key', '')
)[:C.n_shots]
# TODO: Use all/subsampled demos like original (group/sort by consistency_id for relevance), handling context limits (high value, med complexity).
@@ -139,6 +152,49 @@ async def predict_label(example_uid, current_demos, config=C, verbose=False):
# if verbose>1:
# instruction = "Examine the examples below and determine the category 0 or 1 for the final claim based on the pattern. After
# Randomize order to avoid positional bias
random.shuffle(relevant_demos)
messages = [
{"role": "system", "content": instruction}
]
for demo in relevant_demos:
messages.append({"role": "user", "content": demo['prompt']})
messages.append({"role": "assistant", "content": str(demo['label'])})
target_demo = full_demos[example_uid]
messages.append({"role": "user", "content": target_demo['prompt']})
if verbose:
logger.info(f"Predicting for UID {example_uid} with {len(relevant_demos)} shots:")
logger.info(print_messages(messages[:3])) # Log first few for brevity
if verbose > 1:
logger.info(print_messages(messages))
try:
completion = await cached_openrouter_completion_wlogprobs(
model_id=config.model_id,
messages=messages,
max_completion_tokens=60 if verbose else 4,
temperature=0.0,
logprobs=True,
top_logprobs=5,
provider_whitelist=config.provider_whitelist if config.provider_whitelist else None
)
# Extract logprobs for '0' and '1'
choice_logp_dict, logp_dict = get_logprobs_choices(completion, choices=['A', 'B'], regex='[AB]')
if verbose:
logger.info(f"Predicted: {pred}, logprob: {lprob:.4f}")
return pred, lprob
except Exception as e:
logger.error(f"Error in predict_label for UID {example_uid}: {e}")
return random.choice([0, 1]), 0.0
# %% [code]
# Compute energy and metrics
def compute_energy(demos, config=C):
@@ -188,7 +244,7 @@ def compute_energy(demos, config=C):
logger.info("Initial energy: {}", compute_energy(demonstrations))
# %% [code]
def fix_inconsistencies_simple(demos, config=C, max_fixes=5):
async def fix_inconsistencies_simple(demos, config=C, max_fixes=5):
"""Simple consistency fix: for inconsistent pairs, enumerate label combos, re-predict, pick max energy."""
for fix_iter in range(max_fixes):
# Find inconsistent pairs
@@ -226,8 +282,8 @@ def fix_inconsistencies_simple(demos, config=C, max_fixes=5):
# Re-predict labels with new context to update scores
current_labeled = {k: v for k, v in temp_demos.items() if v['label'] is not None}
new_label1, score1 = predict_label(uid1, current_labeled, config)
new_label2, score2 = predict_label(uid2, current_labeled, config)
new_label1, score1 = await predict_label(uid1, temp_demos, current_labeled, config)
new_label2, score2 = await predict_label(uid2, temp_demos, current_labeled, config)
temp_demos[uid1]['score'] = score1
temp_demos[uid2]['score'] = score2
@@ -252,11 +308,20 @@ async def run_icm(demonstrations, config=C):
accuracies = []
# Fix any initial inconsistencies from random initialization
demonstrations = fix_inconsistencies_simple(demonstrations, config)
demonstrations = await fix_inconsistencies_simple(demonstrations, config)
current_labeled = {k: v for k, v in demonstrations.items() if v['label'] is not None}
old_energy, old_metrics = compute_energy(demonstrations, config)
# Set initial scores for seeds by predicting (using current labels)
initial_labeled = {k: v for k, v in demonstrations.items() if v['label'] is not None}
for uid in initial_labeled:
pred_label, score = await predict_label(uid, demonstrations, initial_labeled, config, verbose=False)
demonstrations[uid]['score'] = score # Keep original label, update score
initial_labeled[uid]['score'] = score
old_energy, old_metrics = compute_energy(demonstrations, config)
logger.info("Initial scores set. Updated energy: {}", old_energy)
for iter in range(config.max_iters):
# Weighted sampling for inconsistent groups
groups = {}
@@ -307,13 +372,17 @@ async def run_icm(demonstrations, config=C):
verbose = 2
else:
verbose = 0
new_label, score = predict_label(example_uid, current_labeled, config, verbose=verbose)
new_label, score = await predict_label(example_uid, demonstrations, current_labeled, config, verbose=verbose)
# Update with new label and fix any inconsistencies
temp_demos = demonstrations.copy()
temp_demos[example_uid]['label'] = new_label
temp_demos[example_uid]['score'] = score
temp_demos = fix_inconsistencies_simple(temp_demos, config)
temp_demos = await fix_inconsistencies_simple(temp_demos, config)
# TODO check if we should enable following
# if new_label != demonstrations[example_uid].get('label', None):
# current_labeled = {k: v for k, v in temp_demos.items() if v['label'] is not None}
# Compute new energy
new_energy, new_metrics = compute_energy(temp_demos, config)
pyproject.toml
+1
View File
@@ -9,6 +9,7 @@ authors = [
requires-python = ">=3.10"
dependencies = [
"adjusttext>=1.3.0",
"aiocache>=0.12.3",
"alembic>=1.16.5",
"altair>=5.5.0",
"anthropic>=0.69.0",
uv.lock
Generated
+11
View File
@@ -26,6 +26,15 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/53/1c/8feedd607cc14c5df9aef74fe3af9a99bf660743b842a9b5b1865326b4aa/adjustText-1.3.0-py3-none-any.whl", hash = "sha256:da23d7b24b6db5ffa039bb136bfa556207365e32f48ac74b07ad26dd485bc691", size = 13154, upload-time = "2024-10-31T16:45:35.227Z" },
]
[[package]]
name = "aiocache"
version = "0.12.3"
source = { registry = "https://pypi.org/simple" }
sdist = { url = "https://files.pythonhosted.org/packages/7a/64/b945b8025a9d1e6e2138845f4022165d3b337f55f50984fbc6a4c0a1e355/aiocache-0.12.3.tar.gz", hash = "sha256:f528b27bf4d436b497a1d0d1a8f59a542c153ab1e37c3621713cb376d44c4713", size = 132196, upload-time = "2024-09-25T13:20:23.823Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/37/d7/15d67e05b235d1ed8c3ce61688fe4d84130e72af1657acadfaac3479f4cf/aiocache-0.12.3-py2.py3-none-any.whl", hash = "sha256:889086fc24710f431937b87ad3720a289f7fc31c4fd8b68e9f918b9bacd8270d", size = 28199, upload-time = "2024-09-25T13:20:22.688Z" },
]
[[package]]
name = "aiohappyeyeballs"
version = "2.6.1"
@@ -3748,6 +3757,7 @@ version = "0.1.0"
source = { editable = "." }
dependencies = [
{ name = "adjusttext" },
{ name = "aiocache" },
{ name = "alembic" },
{ name = "altair" },
{ name = "anthropic" },
@@ -3793,6 +3803,7 @@ dev = [
[package.metadata]
requires-dist = [
{ name = "adjusttext", specifier = ">=1.3.0" },
{ name = "aiocache", specifier = ">=0.12.3" },
{ name = "alembic", specifier = ">=1.16.5" },
{ name = "altair", specifier = ">=5.5.0" },
{ name = "anthropic", specifier = ">=0.69.0" },