diff --git a/axbench/__init__.py b/axbench/__init__.py index 7c2515e9..155b591d 100644 --- a/axbench/__init__.py +++ b/axbench/__init__.py @@ -34,6 +34,8 @@ from .models.preference_vector import * from .models.concept_vector import * from .models.hypersteer import * +from .models.latentqa import * +from .models.activation_oracle import * from .models.hypernet.configuration_hypernet import * from .models.hypernet.layers import * diff --git a/axbench/data/prepare_data.py b/axbench/data/prepare_data.py new file mode 100644 index 00000000..e2e5dc78 --- /dev/null +++ b/axbench/data/prepare_data.py @@ -0,0 +1,102 @@ +"""Download pyvene/axbench-concept500 from HuggingFace and prepare it +for concept detection inference (LatentQA, Activation Oracles, PromptDetection). + +Creates the directory structure expected by inference.py: + {dump_dir}/generate/metadata.jsonl + {dump_dir}/train/config.json + {dump_dir}/inference/latent_eval_data.parquet (for overwrite_inference_data_dir) + +IMPORTANT: metadata.jsonl is built from the parquet's own concept descriptions +(output_concept column), NOT from the neuronpedia JSON. The HF dataset contains +a specific curated set of 500 concepts whose IDs don't correspond to the first +500 entries of the neuronpedia JSON. +""" +import os, json, argparse +import pandas as pd +from huggingface_hub import hf_hub_download + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--dump_dir", type=str, required=True) + parser.add_argument("--hf_subdir", type=str, default="2b/l20", + help="Subdirectory in pyvene/axbench-concept500 (e.g. 2b/l20, 9b/l31)") + parser.add_argument("--layer", type=int, default=20, + help="Layer number for train/config.json") + parser.add_argument("--component", type=str, default="res", + help="Component for train/config.json") + args = parser.parse_args() + + dump_dir = args.dump_dir + os.makedirs(f"{dump_dir}/generate", exist_ok=True) + os.makedirs(f"{dump_dir}/train", exist_ok=True) + os.makedirs(f"{dump_dir}/inference", exist_ok=True) + + # 1. Download parquet files from HF + print(f"Downloading parquet files from pyvene/axbench-concept500/{args.hf_subdir}...") + test_path = hf_hub_download( + repo_id="pyvene/axbench-concept500", + filename=f"{args.hf_subdir}/test/data.parquet", + repo_type="dataset", + ) + train_path = hf_hub_download( + repo_id="pyvene/axbench-concept500", + filename=f"{args.hf_subdir}/train/data.parquet", + repo_type="dataset", + ) + test_df = pd.read_parquet(test_path) + train_df = pd.read_parquet(train_path) + print(f"Loaded test: {len(test_df)} rows, train: {len(train_df)} rows") + + # 2. Build metadata.jsonl from the parquet's actual concepts + # This ensures concept_id -> concept description alignment + print("Building metadata.jsonl from parquet concepts...") + concept_map = {} + for _, row in test_df.drop_duplicates("concept_id").iterrows(): + cid = int(row["concept_id"]) + concept_map[cid] = { + "concept": row["output_concept"], + "genre": row.get("concept_genre", "text"), + "sae_link": row.get("sae_link", ""), + } + + metadata_path = f"{dump_dir}/generate/metadata.jsonl" + with open(metadata_path, "w") as f: + for concept_id in sorted(concept_map.keys()): + info = concept_map[concept_id] + entry = { + "concept_id": concept_id, + "concept": info["concept"], + "ref": info["sae_link"], + "concept_genres_map": {info["concept"]: [info["genre"]]}, + } + f.write(json.dumps(entry) + "\n") + print(f"Wrote {len(concept_map)} concepts to {metadata_path}") + + # 3. Create config.json + config = {"layer": args.layer, "component": args.component} + config_path = f"{dump_dir}/train/config.json" + with open(config_path, "w") as f: + json.dump(config, f) + print(f"Wrote config to {config_path}") + + # 4. Save test split as latent_eval_data.parquet + eval_df = test_df.copy() + if "sae_id" not in eval_df.columns: + eval_df["sae_id"] = eval_df["concept_id"] + eval_path = f"{dump_dir}/inference/latent_eval_data.parquet" + eval_df.to_parquet(eval_path, index=False) + print(f"Wrote {len(eval_df)} rows to {eval_path}") + + # 5. Save train data + if "sae_id" not in train_df.columns: + train_df["sae_id"] = train_df["concept_id"] + train_out = f"{dump_dir}/generate/train_data.parquet" + train_df.to_parquet(train_out, index=False) + print(f"Wrote {len(train_df)} rows to {train_out}") + + print("Data preparation complete!") + + +if __name__ == "__main__": + main() diff --git a/axbench/data/setup-latentqa.sh b/axbench/data/setup-latentqa.sh new file mode 100755 index 00000000..018fc748 --- /dev/null +++ b/axbench/data/setup-latentqa.sh @@ -0,0 +1,93 @@ +#!/bin/bash +# Setup script for LatentQA integration with AxBench. +# +# This script: +# 1. Clones and installs the LatentQA repository +# 2. Downloads the pre-trained LatentQA decoder model +# 3. Verifies the installation +# +# Prerequisites: +# - Python 3.10+ with PyTorch and CUDA +# - At least 2 GPUs (target model on cuda:0, decoder on cuda:1) +# - HuggingFace access to meta-llama/Meta-Llama-3-8B-Instruct +# +# Usage: +# bash axbench/data/setup-latentqa.sh [--latentqa-dir /path/to/install] +# +set -e + +LATENTQA_DIR="${1:-./latentqa}" + +echo "=== LatentQA Setup for AxBench ===" +echo "" + +# Step 1: Clone LatentQA repo +if [ -d "$LATENTQA_DIR" ]; then + echo "[1/4] LatentQA repo already exists at $LATENTQA_DIR, pulling latest..." + cd "$LATENTQA_DIR" && git pull && cd - +else + echo "[1/4] Cloning LatentQA repo to $LATENTQA_DIR..." + git clone https://github.com/aypan17/latentqa.git "$LATENTQA_DIR" +fi + +# Step 2: Install LatentQA +echo "[2/4] Installing LatentQA..." +pip install -e "$LATENTQA_DIR" +# Install additional dependencies if requirements.txt exists +if [ -f "$LATENTQA_DIR/requirements.txt" ]; then + pip install -r "$LATENTQA_DIR/requirements.txt" +fi + +# Step 3: Pre-download the decoder model from HuggingFace +echo "[3/4] Pre-downloading LatentQA decoder model (aypan17/latentqa_llama-3-8b-instruct)..." +python3 -c " +from huggingface_hub import snapshot_download +snapshot_download('aypan17/latentqa_llama-3-8b-instruct') +print('Decoder model downloaded successfully.') +" || echo "WARNING: Could not pre-download decoder model. It will be downloaded on first use." + +# Step 4: Verify installation +echo "[4/4] Verifying installation..." +python3 -c " +import torch +print(f'PyTorch version: {torch.__version__}') +print(f'CUDA available: {torch.cuda.is_available()}') +print(f'Number of GPUs: {torch.cuda.device_count()}') +if torch.cuda.device_count() < 2: + print('WARNING: LatentQA requires at least 2 GPUs (target + decoder)') + +# Check LatentQA imports +try: + from lit.utils.activation_utils import latent_qa + from lit.utils.infra_utils import get_model, get_tokenizer, get_modules + from lit.utils.dataset_utils import BASE_DIALOG + print('LatentQA imports: OK') +except ImportError as e: + print(f'LatentQA imports: FAILED ({e})') + print('Make sure the LatentQA repo is on your PYTHONPATH or installed with pip install -e') + exit(1) + +# Check AxBench LatentQA integration +try: + from axbench.models.latentqa import LatentQAReading, LatentQASteering + print('AxBench LatentQA integration: OK') +except ImportError as e: + print(f'AxBench LatentQA integration: FAILED ({e})') + exit(1) + +print() +print('=== Setup complete! ===') +print() +print('To run LatentQA reading mode (concept detection):') +print(' torchrun --nproc_per_node=1 axbench/scripts/inference.py \\\\') +print(' --config axbench/sweep/aryaman/latentqa/reading_llama3_8b.yaml --mode latent') +print() +print('To run LatentQA steering mode:') +print(' torchrun --nproc_per_node=1 axbench/scripts/train.py \\\\') +print(' --config axbench/sweep/aryaman/latentqa/steering_llama3_8b.yaml') +print(' torchrun --nproc_per_node=1 axbench/scripts/inference.py \\\\') +print(' --config axbench/sweep/aryaman/latentqa/steering_llama3_8b.yaml --mode steering') +" + +echo "" +echo "Done!" diff --git a/axbench/models/activation_oracle.py b/axbench/models/activation_oracle.py new file mode 100644 index 00000000..3f089edc --- /dev/null +++ b/axbench/models/activation_oracle.py @@ -0,0 +1,467 @@ +""" +Activation Oracle integration for AxBench. + +Implements concept detection using activation oracles — LoRA-finetuned LLMs +that interpret activations from a target model by receiving them as additive +steering vectors injected at an early layer. + +Based on: https://github.com/adamkarvonen/activation_oracles +Paper: https://arxiv.org/abs/2512.15674 + +Requires the activation_oracles repo: + pip install activation_oracles + Or: pip install 'axbench[activation_oracles]' +""" +import os +import re +import sys + +import torch +from tqdm.auto import tqdm + +from .model import BaseModel + +import logging +logging.basicConfig( + format='%(asctime)s,%(msecs)03d %(levelname)-8s [%(filename)s:%(lineno)d] %(message)s', + datefmt='%Y-%m-%d:%H:%M:%S', + level=logging.WARN, +) +logger = logging.getLogger(__name__) + + +# ── Optional dependency: activation_oracles (nl_probes) ────────────────────── +# Add local clone to sys.path if present: axbench/models/_activation_oracles/ +_ao_local = os.path.join(os.path.dirname(__file__), "_activation_oracles") +if os.path.isdir(_ao_local) and _ao_local not in sys.path: + sys.path.insert(0, _ao_local) + +try: + from nl_probes.utils.activation_utils import collect_activations, get_hf_submodule + from nl_probes.utils.steering_hooks import get_hf_activation_steering_hook, add_hook + from nl_probes.utils.dataset_utils import ( + SPECIAL_TOKEN, get_introspection_prefix, find_pattern_in_tokens, + ) + from nl_probes.utils.common import layer_percent_to_layer + _HAS_AO = True +except ImportError: + _HAS_AO = False + + +def _require_ao(): + if not _HAS_AO: + raise ImportError( + "activation_oracles is not installed. Clone into axbench/models/:\n" + " cd axbench/models && git clone https://github.com/adamkarvonen/activation_oracles.git _activation_oracles\n" + "Or add the repo to PYTHONPATH." + ) + + +# ── Prompt templates ───────────────────────────────────────────────────────── +CONCEPT_DETECTION_QUESTION_TEMPLATE_RATING = ( + "Is the following concept present in the text? Concept: '{concept}'. " + "Answer with a rating from 0 to 2, where 0 means the concept is not present at all, " + "1 means the concept is somewhat present, and 2 means the concept is strongly present. " + "Provide your rating using this exact format: Rating: [[score]]." +) + +CONCEPT_DETECTION_QUESTION_TEMPLATE = ( + "Is the assistant thinking about the concept '{concept}'?" +) + + +# ── Model class ────────────────────────────────────────────────────────────── + +class ActivationOracleReading(BaseModel): + """Activation Oracle concept-detection for AxBench. + + **No training required.** Uses a pre-trained oracle LoRA adapter + (e.g. ``adamkarvonen/checkpoints_latentqa_cls_past_lens_addition_Qwen3-8B``) + loaded onto the same base model as the target. + + Flow for each example: + + 1. Run the *target model* (base weights, oracle LoRA disabled) on the + input text and extract residual-stream activations at + ``layer_percent`` depth. + 2. Build an introspection prompt whose ``" ?"`` placeholder tokens mark + where activations will be injected. + 3. Switch to the *oracle adapter*, register a steering hook at + ``injection_layer`` (default 1), and generate a response. + 4. Parse the response for a 0-2 concept-presence rating. + + Constructor kwargs (beyond the standard ``model, tokenizer, layer``): + + ==================== ==================================================== + ``oracle_lora_path`` HuggingFace repo or local path to the oracle LoRA. + ``target_model_name`` Model name (for ``apply_chat_template``). + ``layer_percent`` Activation extraction depth as % (default 50). + ``injection_layer`` Layer in the oracle to inject at (default 1). + ``steering_coefficient`` Injection strength (default 1.0). + ``max_new_tokens`` Max tokens to generate per question (default 50). + ==================== ==================================================== + """ + + def __init__(self, model, tokenizer, layer=15, training_args=None, **kwargs): + _require_ao() + self.model = model + self.tokenizer = tokenizer + self.layer = layer + self.training_args = training_args + self.device = kwargs.get("device", "cuda:0") + self.seed = kwargs.get("seed", 42) + + # Oracle config — auto-detect LoRA path from model name if not provided + _ORACLE_LORA_MAP = { + "meta-llama/Llama-3.1-8B-Instruct": "adamkarvonen/checkpoints_latentqa_cls_past_lens_Llama-3_1-8B-Instruct", + "Qwen/Qwen3-8B": "adamkarvonen/checkpoints_latentqa_cls_past_lens_addition_Qwen3-8B", + } + model_name_str = getattr(tokenizer, "name_or_path", "") + default_lora = _ORACLE_LORA_MAP.get( + model_name_str, + "adamkarvonen/checkpoints_latentqa_cls_past_lens_addition_Qwen3-8B", + ) + self.oracle_lora_path = kwargs.get("oracle_lora_path", default_lora) + self.target_model_name = kwargs.get( + "target_model_name", model_name_str or "Qwen/Qwen3-8B") + self.layer_percent = kwargs.get("layer_percent", 50) + self.injection_layer = kwargs.get("injection_layer", 1) + self.steering_coefficient = kwargs.get("steering_coefficient", 1.0) + self.max_new_tokens = kwargs.get("max_new_tokens", 50) + + # Internal state (populated by load()) + self._extraction_layer = None + self._oracle_adapter_name = None + self._original_model = model + + # ── BaseModel interface ────────────────────────────────────────────── + + def __str__(self): + return "ActivationOracleReading" + + def make_model(self, **kwargs): + pass + + def save(self, dump_dir, **kwargs): + pass + + def train(self, examples, **kwargs): + pass + + def load(self, dump_dir=None, **kwargs): + """Load the oracle LoRA adapter onto the target model.""" + if self._oracle_adapter_name is not None: + return # already loaded + + from peft import PeftModel, LoraConfig + + # Compute absolute extraction layer + self._extraction_layer = layer_percent_to_layer( + self.target_model_name, self.layer_percent) + logger.warning( + f"Extraction layer: {self._extraction_layer} " + f"({self.layer_percent}% of model depth)") + + # Make model a PeftModel if it isn't already + if not isinstance(self.model, PeftModel): + dummy_config = LoraConfig( + r=1, lora_alpha=1, target_modules=["q_proj"], bias="none", + ) + self.model = PeftModel( + self.model, dummy_config, adapter_name="base_default") + + # Load the oracle LoRA adapter + adapter_name = "activation_oracle" + self.model.load_adapter(self.oracle_lora_path, adapter_name=adapter_name) + self._oracle_adapter_name = adapter_name + logger.warning( + f"Loaded oracle adapter '{adapter_name}' from " + f"{self.oracle_lora_path}") + + # Tokenizer setup for batched generation + self.tokenizer.padding_side = "left" + if self.tokenizer.pad_token_id is None: + self.tokenizer.pad_token_id = self.tokenizer.eos_token_id + + # Pre-compute Yes/No token ids + self._yes_token_id = self.tokenizer.encode( + "Yes", add_special_tokens=False)[0] + self._no_token_id = self.tokenizer.encode( + "No", add_special_tokens=False)[0] + logger.warning( + f"Yes token: {self._yes_token_id}, No token: {self._no_token_id}") + + # ── Core inference ─────────────────────────────────────────────────── + + @torch.no_grad() + def predict_latent(self, examples, **kwargs): + """Run activation-oracle concept detection. + + Returns ``{"max_act": [rating_per_example]}``. + """ + self.model.eval() + concept = kwargs.get("concept", "") + batch_size = kwargs.get("batch_size", 4) + + question_text = CONCEPT_DETECTION_QUESTION_TEMPLATE.format( + concept=concept) + + all_max_act = [] + use_lora = isinstance(self.model, __import__('peft').PeftModel) + + for i in tqdm( + range(0, len(examples), batch_size), + desc="ActivationOracle Reading", + ): + batch_examples = examples.iloc[i : i + batch_size] + + # ── 1. Extract activations from target (base) model ────── + texts = [] + for _, row in batch_examples.iterrows(): + texts.append(row.get("output", row.get("input", ""))) + + self.model.disable_adapter_layers() + + target_inputs = self.tokenizer( + texts, return_tensors="pt", padding=True, truncation=True, + ).to(self.device) + + extraction_submodule = get_hf_submodule( + self.model, self._extraction_layer, use_lora=use_lora) + activations = collect_activations( + self.model, extraction_submodule, target_inputs) + + self.model.enable_adapter_layers() + self.model.set_adapter(self._oracle_adapter_name) + + # ── 2. Build oracle prompts with introspection prefix ──── + oracle_prompts = [] + per_example_vecs = [] + + for b in range(len(texts)): + if target_inputs.attention_mask is not None: + num_valid = int( + target_inputs.attention_mask[b].sum().item()) + else: + num_valid = target_inputs.input_ids.shape[1] + + act_vecs = activations[b, -num_valid:] + per_example_vecs.append(act_vecs) + + prefix = get_introspection_prefix( + self._extraction_layer, num_valid) + user_msg = prefix + question_text + + prompt_str = self.tokenizer.apply_chat_template( + [{"role": "user", "content": user_msg}], + tokenize=False, add_generation_prompt=True, + ) + oracle_prompts.append(prompt_str) + + oracle_inputs = self.tokenizer( + oracle_prompts, return_tensors="pt", + padding=True, truncation=True, + ).to(self.device) + + # ── 3. Locate " ?" positions & build steering vectors ──── + all_positions = [] + all_vectors = [] + + for b in range(oracle_inputs.input_ids.shape[0]): + num_valid = per_example_vecs[b].shape[0] + positions = find_pattern_in_tokens( + oracle_inputs.input_ids[b].tolist(), + SPECIAL_TOKEN, num_valid, self.tokenizer) + vecs = per_example_vecs[b] + n = min(len(positions), vecs.shape[0]) + all_positions.append(positions[:n]) + all_vectors.append(vecs[:n]) + + # ── 4. Forward pass with steering hook at injection layer ── + injection_submodule = get_hf_submodule( + self.model, self.injection_layer, use_lora=use_lora) + hook_fn = get_hf_activation_steering_hook( + all_vectors, all_positions, + self.steering_coefficient, + device=self.device, + dtype=activations.dtype, + ) + + with add_hook(injection_submodule, hook_fn): + outputs = self.model(**oracle_inputs) + + self.model.set_adapter("base_default") + + # ── 5. Extract Yes/No logits at last position ───────── + logits = outputs.logits + attn_mask = oracle_inputs.attention_mask + last_pos = attn_mask.sum(dim=1) - 1 + + for b in range(logits.shape[0]): + last_logits = logits[b, last_pos[b]] + yes_logit = last_logits[self._yes_token_id].item() + no_logit = last_logits[self._no_token_id].item() + score = torch.softmax( + torch.tensor([yes_logit, no_logit]), dim=0)[0].item() + all_max_act.append(score) + + torch.cuda.empty_cache() + + return {"max_act": all_max_act} + + def predict_latents(self, examples, **kwargs): + return self.predict_latent(examples, **kwargs) + + def pre_compute_mean_activations(self, dump_dir, **kwargs): + return {} + + def to(self, device): + self.device = device + return self + + +class ActivationOracleReadingRating(ActivationOracleReading): + """Activation Oracle using 0-2 rating generation instead of Yes/No logits.""" + + def __str__(self): + return "ActivationOracleReadingRating" + + @staticmethod + def _get_rating_from_completion(completion): + """Parse a 0-2 rating from the oracle's free-text completion.""" + try: + if "Rating:" in completion: + rating_text = completion.split("Rating:")[-1].strip() + rating_text = rating_text.split("\n")[0].strip() + rating_text = ( + rating_text.replace("[", "") + .replace("]", "") + .strip('"') + .strip("'") + .strip("*") + .strip() + ) + rating = float(rating_text) + if 0 <= rating <= 2: + return rating + numbers = re.findall(r"\b([012](?:\.\d+)?)\b", completion) + if numbers: + return float(numbers[-1]) + logger.warning( + f"Cannot find rating in completion: {completion[:200]}") + return -1 + except (ValueError, IndexError) as e: + logger.error( + f"Error parsing rating: {completion[:200]}. Error: {e}") + return -1 + + @torch.no_grad() + def predict_latent(self, examples, **kwargs): + self.model.eval() + concept = kwargs.get("concept", "") + batch_size = kwargs.get("batch_size", 4) + + question_text = CONCEPT_DETECTION_QUESTION_TEMPLATE_RATING.format( + concept=concept) + + all_max_act = [] + use_lora = isinstance(self.model, __import__('peft').PeftModel) + + for i in tqdm( + range(0, len(examples), batch_size), + desc="ActivationOracle Rating", + ): + batch_examples = examples.iloc[i : i + batch_size] + + texts = [] + for _, row in batch_examples.iterrows(): + texts.append(row.get("output", row.get("input", ""))) + + self.model.disable_adapter_layers() + + target_inputs = self.tokenizer( + texts, return_tensors="pt", padding=True, truncation=True, + ).to(self.device) + + extraction_submodule = get_hf_submodule( + self.model, self._extraction_layer, use_lora=use_lora) + activations = collect_activations( + self.model, extraction_submodule, target_inputs) + + self.model.enable_adapter_layers() + self.model.set_adapter(self._oracle_adapter_name) + + oracle_prompts = [] + per_example_vecs = [] + + for b in range(len(texts)): + if target_inputs.attention_mask is not None: + num_valid = int(target_inputs.attention_mask[b].sum().item()) + else: + num_valid = target_inputs.input_ids.shape[1] + + act_vecs = activations[b, -num_valid:] + per_example_vecs.append(act_vecs) + + prefix = get_introspection_prefix( + self._extraction_layer, num_valid) + user_msg = prefix + question_text + + prompt_str = self.tokenizer.apply_chat_template( + [{"role": "user", "content": user_msg}], + tokenize=False, add_generation_prompt=True, + ) + oracle_prompts.append(prompt_str) + + oracle_inputs = self.tokenizer( + oracle_prompts, return_tensors="pt", padding=True, truncation=True, + ).to(self.device) + + all_positions = [] + all_vectors = [] + + for b in range(oracle_inputs.input_ids.shape[0]): + num_valid = per_example_vecs[b].shape[0] + positions = find_pattern_in_tokens( + oracle_inputs.input_ids[b].tolist(), + SPECIAL_TOKEN, num_valid, self.tokenizer) + vecs = per_example_vecs[b] + n = min(len(positions), vecs.shape[0]) + all_positions.append(positions[:n]) + all_vectors.append(vecs[:n]) + + injection_submodule = get_hf_submodule( + self.model, self.injection_layer, use_lora=use_lora) + hook_fn = get_hf_activation_steering_hook( + all_vectors, all_positions, + self.steering_coefficient, + device=self.device, + dtype=activations.dtype, + ) + + with add_hook(injection_submodule, hook_fn): + outputs = self.model.generate( + **oracle_inputs, + max_new_tokens=self.max_new_tokens, + do_sample=False, + ) + + self.model.set_adapter("base_default") + + prompt_len = oracle_inputs.input_ids.shape[1] + for b in range(outputs.shape[0]): + completion_ids = outputs[b, prompt_len:] + completion = self.tokenizer.decode( + completion_ids, skip_special_tokens=True) + rating = self._get_rating_from_completion(completion) + if len(all_max_act) < 20: + logger.warning( + f"[AO Rating] text={texts[b][:80]}... " + f"completion={completion!r} rating={rating}") + all_max_act.append(rating) + + torch.cuda.empty_cache() + + return {"max_act": all_max_act} + + diff --git a/axbench/models/latentqa.py b/axbench/models/latentqa.py new file mode 100644 index 00000000..5d1cc028 --- /dev/null +++ b/axbench/models/latentqa.py @@ -0,0 +1,436 @@ +""" +LatentQA integration for AxBench (reading only). + +Implements two model classes: +- LatentQAReading: Uses LatentQA's reading mode for concept detection (latent inference). +- LatentQAReadingRating: Variant that uses free-text rating generation instead of Yes/No logits. + +Requires the LatentQA repo (https://github.com/aypan17/latentqa) to be installed. +""" +import json +import os +import sys +import re + +import numpy as np +import torch +from tqdm.auto import tqdm + +from transformers import AutoModelForCausalLM +from .model import Model, BaseModel + +import logging +logging.basicConfig( + format='%(asctime)s,%(msecs)03d %(levelname)-8s [%(filename)s:%(lineno)d] %(message)s', + datefmt='%Y-%m-%d:%H:%M:%S', + level=logging.WARN, +) +logger = logging.getLogger(__name__) + + +# Yes/No question template for concept detection. +# We measure logits for "Yes" vs "No" tokens from the decoder. +CONCEPT_DETECTION_QUESTION_TEMPLATE = ( + "Is the assistant thinking about the concept '{concept}'?" +) + +CONCEPT_DETECTION_QUESTION_TEMPLATE_RATING = ( + "Is the following concept present in the text? Concept: '{concept}'. " + "Answer with a rating from 0 to 2, where 0 means the concept is not present at all, " + "1 means the concept is somewhat present, and 2 means the concept is strongly present. " + "Provide your rating using this exact format: Rating: [[score]]." +) + + +# Add local clone to sys.path if present: axbench/models/_latentqa/ +_latentqa_local = os.path.join(os.path.dirname(__file__), "_latentqa") +if os.path.isdir(_latentqa_local) and _latentqa_local not in sys.path: + sys.path.insert(0, _latentqa_local) + +try: + from lit.utils.activation_utils import latent_qa as _latent_qa + from lit.utils.dataset_utils import BASE_DIALOG as _BASE_DIALOG, ENCODER_CHAT_TEMPLATES as _ENCODER_CHAT_TEMPLATES + from lit.utils.infra_utils import get_tokenizer as _lqa_get_tokenizer + try: + from lit.utils.dataset_utils import lqa_tokenize as _lqa_tokenize + except ImportError: + from lit.utils.dataset_utils import tokenize as _lqa_tokenize + _HAS_LATENTQA = True +except ImportError: + _HAS_LATENTQA = False + + +def _require_latentqa(): + if not _HAS_LATENTQA: + raise ImportError( + "LatentQA is not installed. Clone into axbench/models/:\n" + " cd axbench/models && git clone https://github.com/aypan17/latentqa.git _latentqa\n" + "Or add the repo to PYTHONPATH." + ) + + +def _get_model_layers_str(model): + """Determine the correct attribute path to model layers.""" + for path in [ + "model.layers", + "model.model.layers", + "module.model.model.layers", + "language_model.model.layers", + "module.language_model.model.layers", + ]: + obj = model + try: + for attr in path.split("."): + obj = getattr(obj, attr) + return path + except AttributeError: + continue + raise RuntimeError("Cannot find model layers. Unsupported model architecture.") + + +def _load_decoder_model(target_model_name, decoder_model_name, decoder_device): + """Load the LatentQA decoder model (shared across model classes). + + Replicates the essential steps of latentqa's ``get_model`` but uses + ``sdpa`` attention so we don't depend on ``flash-attn``. + """ + _require_latentqa() + from peft import PeftModel + + logger.warning(f"Loading LatentQA decoder from {decoder_model_name} to {decoder_device}") + lqa_tokenizer = _lqa_get_tokenizer(target_model_name) + + # Load base model with sdpa (no flash-attn dependency) + base_model = AutoModelForCausalLM.from_pretrained( + target_model_name, + attn_implementation="sdpa", + torch_dtype=torch.bfloat16, + device_map="auto" if decoder_device == "auto" else None, + ) + base_model.resize_token_embeddings(len(lqa_tokenizer)) + for p in base_model.parameters(): + p.requires_grad = False + + # Load LoRA decoder weights + decoder_model = PeftModel.from_pretrained(base_model, decoder_model_name) + if decoder_device is not None and decoder_device != "auto": + decoder_model = decoder_model.to(decoder_device) + decoder_model.eval() + return decoder_model + + +def _get_modules(target_model, decoder_model, min_layer=15, max_layer=16, + layer_to_write=0, num_layers_to_read=1): + """Get read/write module hooks for LatentQA. + + Returns List[List[Module]] for read and write modules. + """ + target_path = _get_model_layers_str(target_model) + decoder_path = _get_model_layers_str(decoder_model) + + def get_layer(model, path, idx): + obj = model + for attr in path.split("."): + obj = getattr(obj, attr) + return obj[idx] + + module_read, module_write = [], [] + for i in range(min_layer, max_layer): + module_read_i = [get_layer(target_model, target_path, j) + for j in range(i, i + num_layers_to_read)] + module_write_i = [get_layer(decoder_model, decoder_path, j) + for j in range(layer_to_write, layer_to_write + num_layers_to_read)] + module_read.append(module_read_i) + module_write.append(module_write_i) + return module_read, module_write + + +class LatentQAReading(BaseModel): + """LatentQA Reading mode for concept detection. + + Uses LatentQA's decoder to read and interpret activations from the + target model, then determines if a concept is present. + + This is similar to PromptDetection but reads from internal activations + rather than prompting the model directly about its output. + """ + + def __init__(self, model, tokenizer, layer=15, training_args=None, **kwargs): + self.model = model # target model + self.tokenizer = tokenizer + self.layer = layer + self.training_args = training_args + self.device = kwargs.get("device", "cuda:0") + self.decoder_device = kwargs.get("decoder_device", "cuda:1") + self.seed = kwargs.get("seed", 42) + + # LatentQA-specific config + self.decoder_model_name = kwargs.get( + "decoder_model_name", "aypan17/latentqa_llama-3-8b-instruct") + self.target_model_name = kwargs.get( + "target_model_name", "meta-llama/Meta-Llama-3-8B-Instruct") + self.min_layer_to_read = kwargs.get("min_layer_to_read", 15) + self.max_layer_to_read = kwargs.get("max_layer_to_read", 16) + self.num_layers_to_read = kwargs.get("num_layers_to_read", 1) + self.layer_to_write = kwargs.get("layer_to_write", 0) + self.modify_chat_template = kwargs.get("modify_chat_template", True) + self.max_new_tokens = kwargs.get("max_new_tokens", 100) + + self.decoder_model = None + self.module_read = None + self.module_write = None + + def __str__(self): + return 'LatentQAReading' + + def make_model(self, **kwargs): + pass + + def save(self, dump_dir, **kwargs): + pass # no training needed + + def train(self, examples, **kwargs): + pass # no training needed + + def load(self, dump_dir=None, **kwargs): + """Load the LatentQA decoder model.""" + if self.decoder_model is not None: + return # already loaded + + self.decoder_model = _load_decoder_model( + self.target_model_name, self.decoder_model_name, self.decoder_device) + + # Ensure decoder vocab matches target model (e.g. if PAD token was added) + target_vocab_size = self.model.get_input_embeddings().weight.shape[0] + decoder_vocab_size = self.decoder_model.get_input_embeddings().weight.shape[0] + if target_vocab_size != decoder_vocab_size: + logger.warning( + f"Resizing decoder embeddings from {decoder_vocab_size} to {target_vocab_size}") + self.decoder_model.resize_token_embeddings(target_vocab_size) + + # Set up read/write module hooks + self.module_read, self.module_write = _get_modules( + self.model, self.decoder_model, + min_layer=self.min_layer_to_read, + max_layer=self.max_layer_to_read, + layer_to_write=self.layer_to_write, + num_layers_to_read=self.num_layers_to_read, + ) + + @torch.no_grad() + def predict_latent(self, examples, **kwargs): + """Use LatentQA reading mode for concept detection. + + For each example, extracts activations from the target model, + feeds them to the LatentQA decoder with a yes/no concept question, + and uses P(Yes) - P(No) logit difference as the detection score. + """ + _require_latentqa() + + self.model.eval() + self.decoder_model.eval() + + concept = kwargs.get("concept", "") + batch_size = kwargs.get("batch_size", 4) + + # Get token IDs for "Yes" and "No" + yes_token_id = self.tokenizer.encode("Yes", add_special_tokens=False)[0] + no_token_id = self.tokenizer.encode("No", add_special_tokens=False)[0] + + question_text = CONCEPT_DETECTION_QUESTION_TEMPLATE.format(concept=concept) + chat_template = _ENCODER_CHAT_TEMPLATES.get(self.tokenizer.name_or_path, None) + + all_max_act = [] + + # LatentQA requires left padding + orig_padding_side = self.tokenizer.padding_side + self.tokenizer.padding_side = "left" + + for i in tqdm(range(0, len(examples), batch_size), desc="LatentQA Reading"): + batch_examples = examples.iloc[i:i + batch_size] + + probe_data = [] + for _, row in batch_examples.iterrows(): + user_text = row.get("input", "") + assistant_text = row.get("output", "") + read_prompt = self.tokenizer.apply_chat_template( + [ + {"role": "user", "content": user_text}, + {"role": "assistant", "content": assistant_text}, + ], + tokenize=False, + add_generation_prompt=False, + chat_template=chat_template, + ) + dialog = _BASE_DIALOG + [ + {"role": "user", "content": question_text}, + ] + probe_data.append({ + "read_prompt": read_prompt, + "dialog": dialog, + }) + + # Tokenize with generate=True to include the assistant header, + # then do a forward pass to get logits (not model.generate). + batch_tokenized = _lqa_tokenize( + probe_data, + self.tokenizer, + name=self.target_model_name, + generate=True, + mask_type=None, + mask_all_but_last=True, + modify_chat_template=self.modify_chat_template, + ) + + # Add dummy labels so latent_qa accepts generate=False + input_ids = batch_tokenized["tokenized_write"]["input_ids"] + batch_tokenized["tokenized_write"]["labels"] = input_ids.clone() + + # Forward pass to get logits + out = _latent_qa( + batch_tokenized, + self.model, + self.decoder_model, + self.module_read[0], + self.module_write[0], + self.tokenizer, + shift_position_ids=False, + generate=False, + no_grad=True, + ) + + # Extract logits at the last non-padding position for each example + logits = out.logits # (batch, seq_len, vocab_size) + attention_mask = batch_tokenized["tokenized_write"]["attention_mask"].to(logits.device) + # Last real token position per example + last_pos = attention_mask.sum(dim=1) - 1 # (batch,) + + for j in range(logits.shape[0]): + last_logits = logits[j, last_pos[j]] # (vocab_size,) + yes_logit = last_logits[yes_token_id].item() + no_logit = last_logits[no_token_id].item() + # Score: P(Yes) from softmax over Yes/No + score = torch.softmax( + torch.tensor([yes_logit, no_logit]), dim=0 + )[0].item() + all_max_act.append(score) + + torch.cuda.empty_cache() + + self.tokenizer.padding_side = orig_padding_side + return {"max_act": all_max_act} + + def predict_latents(self, examples, **kwargs): + return self.predict_latent(examples, **kwargs) + + def pre_compute_mean_activations(self, dump_dir, **kwargs): + return {} + + def to(self, device): + self.device = device + # Note: target model device is managed by AxBench infrastructure + # decoder stays on its own device + return self + + +class LatentQAReadingRating(LatentQAReading): + """LatentQA Reading using 0-2 rating generation instead of Yes/No logits.""" + + def __str__(self): + return "LatentQAReadingRating" + + @staticmethod + def _get_rating_from_completion(completion): + """Parse a 0-2 rating from the decoder's free-text completion.""" + import re + try: + if "Rating:" in completion: + rating_text = completion.split("Rating:")[-1].strip() + rating_text = rating_text.split("\n")[0].strip() + rating_text = ( + rating_text.replace("[", "").replace("]", "") + .strip('"').strip("'").strip("*").strip() + ) + rating = float(rating_text) + if 0 <= rating <= 2: + return rating + numbers = re.findall(r"\b([012](?:\.\d+)?)\b", completion) + if numbers: + return float(numbers[-1]) + return -1 + except (ValueError, IndexError): + return -1 + + def predict_latent(self, examples, **kwargs): + _require_latentqa() + + self.model.eval() + self.decoder_model.eval() + + concept = kwargs.get("concept", "") + batch_size = kwargs.get("batch_size", 4) + + question_text = CONCEPT_DETECTION_QUESTION_TEMPLATE_RATING.format(concept=concept) + chat_template = _ENCODER_CHAT_TEMPLATES.get(self.tokenizer.name_or_path, None) + + all_max_act = [] + + orig_padding_side = self.tokenizer.padding_side + self.tokenizer.padding_side = "left" + + for i in tqdm(range(0, len(examples), batch_size), desc="LatentQA Rating"): + batch_examples = examples.iloc[i:i + batch_size] + + probe_data = [] + for _, row in batch_examples.iterrows(): + user_text = row.get("input", "") + assistant_text = row.get("output", "") + read_prompt = self.tokenizer.apply_chat_template( + [ + {"role": "user", "content": user_text}, + {"role": "assistant", "content": assistant_text}, + ], + tokenize=False, + add_generation_prompt=False, + chat_template=chat_template, + ) + dialog = _BASE_DIALOG + [ + {"role": "user", "content": question_text}, + ] + probe_data.append({ + "read_prompt": read_prompt, + "dialog": dialog, + }) + + batch_tokenized = _lqa_tokenize( + probe_data, self.tokenizer, name=self.target_model_name, + generate=True, mask_type=None, mask_all_but_last=True, + modify_chat_template=self.modify_chat_template, + ) + + with torch.no_grad(): + out = _latent_qa( + batch_tokenized, self.model, self.decoder_model, + self.module_read[0], self.module_write[0], self.tokenizer, + shift_position_ids=False, generate=True, + max_new_tokens=50, no_grad=True, + ) + + num_tokens = batch_tokenized["tokenized_write"]["input_ids"][0].shape[0] + for j in range(len(batch_examples)): + completion = self.tokenizer.decode( + out[j][num_tokens:], skip_special_tokens=True) + rating = self._get_rating_from_completion(completion) + if len(all_max_act) < 20: + text = batch_examples.iloc[j].get("output", "")[:80] + logger.warning( + f"[LQA Rating] text={text}... " + f"completion={completion!r} rating={rating}") + all_max_act.append(rating) + + torch.cuda.empty_cache() + + self.tokenizer.padding_side = orig_padding_side + return {"max_act": all_max_act} + + diff --git a/axbench/scripts/evaluate.py b/axbench/scripts/evaluate.py index 0e15c2d7..9f60ab22 100644 --- a/axbench/scripts/evaluate.py +++ b/axbench/scripts/evaluate.py @@ -668,8 +668,32 @@ def eval_latent(args): if concept_id < start_concept_id: continue logger.warning(f"Evaluating concept_id: {concept_id}") - + # Initialize a dictionary for storing evaluation results for this `concept_id` + eval_results = {} + for evaluator_name in args.latent_evaluators: + for model_name in args.models: + if model_name in LATENT_EXCLUDE_MODELS: + continue + evaluator_class = getattr(axbench, evaluator_name) + evaluator = evaluator_class(model_name) + result = evaluator.compute_metrics(current_df) + if evaluator_name not in eval_results: + eval_results[evaluator_name] = {} + eval_results[evaluator_name][model_name] = result + + save_results( + dump_dir, + {"concept_id": concept_id + 1}, + concept_id, + "latent", + eval_results, + ) + + # Generate plots + plot_latent(dump_dir, args.report_to if hasattr(args, 'report_to') else [], + args.wandb_name if hasattr(args, 'wandb_name') else None) + def main(): custom_args = [ diff --git a/axbench/scripts/inference.py b/axbench/scripts/inference.py index db06b9af..ab52c6d2 100644 --- a/axbench/scripts/inference.py +++ b/axbench/scripts/inference.py @@ -39,7 +39,7 @@ CONFIG_FILE = "config.json" METADATA_FILE = "metadata.jsonl" STEERING_WITH_SHARED_MODELS = {"HyperSteer"} -STEERING_EXCLUDE_MODELS = {"IntegratedGradients", "InputXGradients", "PromptDetection", "BoW"} +STEERING_EXCLUDE_MODELS = {"IntegratedGradients", "InputXGradients", "PromptDetection", "BoW", "LatentQAReading", "ActivationOracleReading"} LATENT_EXCLUDE_MODELS = {"PromptSteering", "PromptBaseline", "DiReFT", "LoReFT", "LoRA", "SFT", "HyperSteer"} LATENT_PROMPT_PREFIX = "Generate a random sentence." @@ -409,6 +409,7 @@ def infer_steering(args, rank, world_size, device, logger, training_args, genera training_args=training_args.models[model_name] if model_name not in {"PromptSteering", "GemmaScopeSAE"} else None, # we init with training args as well low_rank_dimension=len(metadata), device=device, steering_layers=steering_layers, + metadata=metadata, ) if model_name in {"PromptSteering", "GemmaScopeSAE"}: lr = 1 @@ -658,8 +659,11 @@ def infer_latent(args, rank, world_size, device, logger, training_args, generate else: current_df = cache_df[(concept_id, dataset_category)] + predict_kwargs = dict(batch_size=args.latent_batch_size, prefix_length=prefix_length) + if model_name in {"PromptDetection", "LatentQAReading", "LatentQAReadingRating", "ActivationOracleReading", "ActivationOracleReadingRating"}: + predict_kwargs["concept"] = metadata[concept_id]["concept"] results = benchmark_model.predict_latent( - current_df, batch_size=args.latent_batch_size, prefix_length=prefix_length + current_df, **predict_kwargs ) # Store the results in current_df for k, v in results.items(): @@ -852,7 +856,7 @@ def infer_latent_imbalance(args, rank, world_size, device, logger, training_args low_rank_dimension=len(metadata), device=device ) - if model_name in {"PromptDetection", "BoW"}: + if model_name in {"PromptDetection", "BoW", "LatentQAReading", "ActivationOracleReading"}: for concept_id in concept_ids: benchmark_model.load( dump_dir=train_dir, sae_path=metadata[0]["ref"], mode="latent", diff --git a/axbench/scripts/train.py b/axbench/scripts/train.py index c00e3363..48c5979e 100644 --- a/axbench/scripts/train.py +++ b/axbench/scripts/train.py @@ -461,7 +461,7 @@ def main(): intervention_positions_dropout=args.models[model_name].intervention_positions_dropout, preference_pairs=args.models[model_name].preference_pairs, ) - if model_name not in {"LoReFT", "LoRA", "SFT", "BoW", "PreferenceLoReFT", "ConceptLoReFT"} and args.use_bf16: + if model_name not in {"LoReFT", "LoRA", "SFT", "BoW", "PreferenceLoReFT", "ConceptLoReFT", "LatentQAReading", "LatentQAReadingRating", "ActivationOracleReading", "ActivationOracleReadingRating"} and args.use_bf16: if isinstance(benchmark_model.ax, list): for ax in benchmark_model.ax: ax.to(torch.bfloat16) @@ -473,6 +473,9 @@ def main(): "exclude_bos": args.models[model_name].exclude_bos, "metadata_path": metadata_path, "use_dpo_loss": args.use_dpo_loss, + "concept": concept, + "concept_id": concept_id, + "dump_dir": str(dump_dir), "logging_metadata": { "concept_id": concept_id, "model_name": model_name, diff --git a/axbench/sweep/aryaman/activation_oracle/reading_llama3_1_8b.yaml b/axbench/sweep/aryaman/activation_oracle/reading_llama3_1_8b.yaml new file mode 100644 index 00000000..cf4dc118 --- /dev/null +++ b/axbench/sweep/aryaman/activation_oracle/reading_llama3_1_8b.yaml @@ -0,0 +1,63 @@ +# Activation Oracle Reading mode (concept detection) on Llama-3.1-8B-Instruct. +# +# Oracle LoRA: adamkarvonen/checkpoints_latentqa_cls_past_lens_Llama-3_1-8B-Instruct +# Trained at layers 25/50/75% depth; we use 50% (layer 16). +# Injection at layer 1, steering coefficient 1.0. + +generate: + lm_model: "gpt-4o-mini" + output_length: 128 + num_of_examples: 144 + concept_path: "axbench/data/gemma-2-2b_20-gemmascope-res-16k.json" + max_concepts: 500 + master_data_dir: "axbench/data" + dataset_category: "instruction" + lm_use_cache: false + seed: 42 + +train: + model_name: "meta-llama/Llama-3.1-8B-Instruct" + layer: 16 # 50% of 32 layers + component: "res" + seed: 42 + use_bf16: true + +inference: + use_bf16: true + models: ["ActivationOracleReading"] + model_name: "meta-llama/Llama-3.1-8B-Instruct" + # latent related params + output_length: 128 + latent_num_of_examples: 36 + latent_batch_size: 4 + # steering related params (not used for reading, but required by schema) + steering_intervention_type: "addition" + steering_model_name: "meta-llama/Llama-3.1-8B-Instruct" + steering_datasets: ["AlpacaEval"] + steering_batch_size: 4 + steering_output_length: 128 + steering_layers: [16] + steering_num_of_examples: 10 + steering_factors: [0.2, 0.4, 0.6, 0.8, 1.0, 1.2, 1.4, 1.6, 1.8, 2.0, 2.5, 3.0, 4.0, 5.0] + # master data dir + master_data_dir: "axbench/data" + seed: 42 + lm_model: "gpt-4o-mini" + temperature: 1.0 + +evaluate: + models: ["ActivationOracleReading"] + latent_evaluators: [ + "AUCROCEvaluator", + "HardNegativeEvaluator", + ] + steering_evaluators: [ + "PerplexityEvaluator", + "LMJudgeEvaluator", + ] + winrate_split_ratio: 0.5 + num_of_workers: 32 + lm_model: "gpt-4o-mini" + run_winrate: false + winrate_baseline: "PromptSteering" + master_data_dir: "axbench/data" diff --git a/axbench/sweep/aryaman/activation_oracle/reading_qwen3_8b.yaml b/axbench/sweep/aryaman/activation_oracle/reading_qwen3_8b.yaml new file mode 100644 index 00000000..cf021f81 --- /dev/null +++ b/axbench/sweep/aryaman/activation_oracle/reading_qwen3_8b.yaml @@ -0,0 +1,86 @@ +# Activation Oracle Reading mode (concept detection) on Qwen3-8B. +# +# This config evaluates activation oracles for concept detection. +# An oracle is the same base LLM fine-tuned with LoRA to interpret +# hidden-state activations injected at an early layer. +# +# Reference: https://github.com/adamkarvonen/activation_oracles +# Paper: https://arxiv.org/abs/2512.15674 +# +# Prerequisites: +# 1. pip install peft bitsandbytes +# 2. Oracle LoRA is downloaded automatically from HuggingFace: +# adamkarvonen/checkpoints_latentqa_cls_past_lens_addition_Qwen3-8B +# 3. Only 1 GPU required (single model with adapter switching). +# +# Usage: +# # Data generation (reuse existing or generate new): +# torchrun --nproc_per_node=1 axbench/scripts/generate.py \ +# --config axbench/sweep/aryaman/activation_oracle/reading_qwen3_8b.yaml +# +# # Inference (concept detection): +# torchrun --nproc_per_node=1 axbench/scripts/inference.py \ +# --config axbench/sweep/aryaman/activation_oracle/reading_qwen3_8b.yaml --mode latent +# +# # Evaluation: +# python axbench/scripts/evaluate.py \ +# --config axbench/sweep/aryaman/activation_oracle/reading_qwen3_8b.yaml + +generate: + lm_model: "gpt-4o-mini" + output_length: 128 + num_of_examples: 144 + concept_path: "axbench/data/gemma-2-2b_10-gemmascope-res-16k.json" + max_concepts: 500 + master_data_dir: "axbench/data" + dataset_category: "instruction" + lm_use_cache: false + seed: 42 + +train: + model_name: "Qwen/Qwen3-8B" + layer: 18 # ~50% of 36 layers + component: "res" + seed: 42 + use_bf16: true + # ActivationOracleReading does not require training. + +inference: + use_bf16: true + models: ["ActivationOracleReading"] + model_name: "Qwen/Qwen3-8B" + # latent related params + output_length: 128 + latent_num_of_examples: 36 + latent_batch_size: 4 # keep small — generation is the bottleneck + # steering related params (not used for reading, but required by schema) + steering_intervention_type: "addition" + steering_model_name: "Qwen/Qwen3-8B" + steering_datasets: ["AlpacaEval"] + steering_batch_size: 4 + steering_output_length: 128 + steering_layers: [18] + steering_num_of_examples: 10 + steering_factors: [0.2, 0.4, 0.6, 0.8, 1.0, 1.2, 1.4, 1.6, 1.8, 2.0, 2.5, 3.0, 4.0, 5.0] + # master data dir + master_data_dir: "axbench/data" + seed: 42 + lm_model: "gpt-4o-mini" + temperature: 1.0 + +evaluate: + models: ["ActivationOracleReading"] + latent_evaluators: [ + "AUCROCEvaluator", + "HardNegativeEvaluator", + ] + steering_evaluators: [ + "PerplexityEvaluator", + "LMJudgeEvaluator", + ] + winrate_split_ratio: 0.5 + num_of_workers: 32 + lm_model: "gpt-4o-mini" + run_winrate: false + winrate_baseline: "PromptSteering" + master_data_dir: "axbench/data" diff --git a/axbench/sweep/aryaman/latentqa/reading_llama3_8b.yaml b/axbench/sweep/aryaman/latentqa/reading_llama3_8b.yaml new file mode 100644 index 00000000..3e4b837a --- /dev/null +++ b/axbench/sweep/aryaman/latentqa/reading_llama3_8b.yaml @@ -0,0 +1,82 @@ +# LatentQA Reading mode (concept detection) on Llama-3-8B-Instruct. +# +# This config evaluates LatentQA's reading mode for concept detection. +# LatentQA decodes the target model's activations at layer 15 using a +# fine-tuned decoder to determine if a concept is present. +# +# Prerequisites: +# 1. Install LatentQA: pip install -e /path/to/latentqa +# 2. Download decoder: aypan17/latentqa_llama-3-8b-instruct (auto from HF) +# 3. Requires 2 GPUs (target on cuda:0, decoder on cuda:1) +# +# Usage: +# # Data generation (reuse existing or generate new): +# torchrun --nproc_per_node=1 axbench/scripts/generate.py \ +# --config axbench/sweep/aryaman/latentqa/reading_llama3_8b.yaml +# +# # Inference (concept detection): +# torchrun --nproc_per_node=1 axbench/scripts/inference.py \ +# --config axbench/sweep/aryaman/latentqa/reading_llama3_8b.yaml --mode latent +# +# # Evaluation: +# python axbench/scripts/evaluate.py \ +# --config axbench/sweep/aryaman/latentqa/reading_llama3_8b.yaml + +generate: + lm_model: "gpt-4o-mini" + output_length: 128 + num_of_examples: 144 + concept_path: "axbench/data/llama3.1-8b_20-llamascope-res-131k.json" + max_concepts: 500 + master_data_dir: "axbench/data" + dataset_category: "instruction" + lm_use_cache: false + seed: 42 + +train: + model_name: "meta-llama/Meta-Llama-3-8B-Instruct" + layer: 15 + component: "res" + seed: 42 + use_bf16: true + # LatentQAReading does not require training. + +inference: + use_bf16: true + models: ["LatentQAReading"] + model_name: "meta-llama/Meta-Llama-3-8B-Instruct" + # latent related params + output_length: 128 + latent_num_of_examples: 36 + latent_batch_size: 4 # small batch size since 2 models are in GPU memory + # steering related params (not used for reading, but required by schema) + steering_intervention_type: "addition" + steering_model_name: "meta-llama/Meta-Llama-3-8B-Instruct" + steering_datasets: ["AlpacaEval"] + steering_batch_size: 4 + steering_output_length: 128 + steering_layers: [15] + steering_num_of_examples: 10 + steering_factors: [0.2, 0.4, 0.6, 0.8, 1.0, 1.2, 1.4, 1.6, 1.8, 2.0, 2.5, 3.0, 4.0, 5.0] + # master data dir + master_data_dir: "axbench/data" + seed: 42 + lm_model: "gpt-4o-mini" + temperature: 1.0 + +evaluate: + models: ["LatentQAReading"] + latent_evaluators: [ + "AUCROCEvaluator", + "HardNegativeEvaluator", + ] + steering_evaluators: [ + "PerplexityEvaluator", + "LMJudgeEvaluator", + ] + winrate_split_ratio: 0.5 + num_of_workers: 32 + lm_model: "gpt-4o-mini" + run_winrate: false + winrate_baseline: "PromptSteering" + master_data_dir: "axbench/data" diff --git a/pyproject.toml b/pyproject.toml index 74ba9e27..12a78154 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -17,13 +17,19 @@ dependencies = [ "scikit-learn>=1.5.2", "seaborn>=0.12.2", "pyreft>=0.0.8", - "peft>=0.13.2", + "peft>=0.16", "jupyter>=1.1.1", "adjusttext>=1.3.0", "altair>=5.5.0", "umap-learn>=0.5.7", + "stanza>=1.11.0", ] +# latentqa and activation_oracles have pinned deps that conflict with axbench. +# Clone into axbench/models/ instead — they are auto-discovered on sys.path: +# cd axbench/models && git clone https://github.com/aypan17/latentqa.git _latentqa +# cd axbench/models && git clone https://github.com/adamkarvonen/activation_oracles.git _activation_oracles + [build-system] requires = ["hatchling"] build-backend = "hatchling.build"