Count samples and tokens in training data.

Create input_ids and assistant_masks for training, where assistant_masks indicates which tokens should have loss computed (1 for assistant to

Note: Will always make all split(s) at the same time. This is so we ensure that train / test splits have no overlap.

Reads seeded SAE hard negatives and builds per-SAE train/test splits for detection. Ensures all entries share the same SAEInfo and match the

Creates a new dataset where each conversation is expanded into multiple rows with incrementally increasing turns. Required for Qwen3 tokeniza

Create training data from collected activations.

Download a specific folder from a Hugging Face repository.

Note: x can be either shape (B, F) or (B, L, F)

Encode message dictionaries into tokenized inputs.

Return (accuracy, items) for a single .tsv dataset file. items is a list of (prompt, response, ground_truth) tuples.

Use GPT-5-mini to identify which sentences match the explanation. Args: batch: Mixed sentences batch with target explanation

Extract learning rate from column name. Returns None if not a loss column.

Filter bars to only include those with allowed labels. Args: names: List of LoRA names labels: List of legend labels mean

Find the most similar features to a target feature using cosine similarity of encoder vectors.

Fix a single JSON file.

Check if response matches ground truth gender.

Generate explanations for all SAE activations using a specific model.

Calculate best-of-N scores for each record, filtered by the best investigator prompt. Args: data: JsonSchema object loaded from JSON

Calculate best-of-N scores for each record, filtered by the best investigator prompt. Args: data: JsonSchema object loaded from JSON

Calculate best-of-N scores for each record, filtered by the best investigator prompt.

Create mask for BOS, EOS, and PAD tokens

Return model kwargs based on model name.

User prompt should be labeled as -100

Get all ContextQASample objects from specified groups. Args: group_names: List of group names (e.g., ["sst2", "ag_news"]) nu

Returns LoRA target pattern for text-only training on VLMs, or None if not a VLM.

Get a mix of pretrain and chat data at a specified ratio. By default, 90% of the data will be pretrain and 10% will be chat. Default datasets:

Load classification results for a single base model and compute OOD accuracy per LoRA checkpoint. Returns dict mapping lora_name -> (mean, er

Load classification results from folder.

Load hardcoded Claude results from TSV data. Args: is_open_ended: If True, return open-ended results; if False, return yes/no results

Load all classification datasets for a specific model and layer percent.

Load Gender results for Gemma-2-9B-IT, matching plot_secret_keeping_results.py (GENDER_SEQUENCE=True and CHOSEN_GENDER_PROMPT). Returns d

Load all JSON files from the directory.

Load knowledge eval results from folder. Args: folder_path: Path to folder containing base_model.json and personaqa_lora.json ver

Non-exhaustive list of models. For demonstration purposes. Simply copy and create a new function for your needs.

Load PersonaQA open-ended sequence results and compute mean sequence-level accuracy per LoRA. Uses the same logic as plot_personaqa_results_a

Load all JSON files from the directory.

Load PersonAQA results from directory and _orig subdirectory.

Load all JSON files from the directory.

Load all JSON files from the directory.

Load all JSON files from the directory.

Load all JSON files from the directory.

Load all JSON files from the directory.