diff --git a/nbs/simple_icm.py b/nbs/simple_icm.py index 479e870..df229ce 100644 --- a/nbs/simple_icm.py +++ b/nbs/simple_icm.py @@ -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) diff --git a/pyproject.toml b/pyproject.toml index 81cba5e..72ce687 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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", diff --git a/uv.lock b/uv.lock index d9a0948..1bae939 100644 --- a/uv.lock +++ b/uv.lock @@ -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" },