persona-vectors 0.8.0

Library for extracting and analyzing persona vectors

Persona Vectors

Extract persona-aligned activation vectors from language models and analyze how persona prompts move hidden states.

This project is experimental.

Install

uv sync
cp .env.example .env

Python >=3.12 is required. Set NDIF_API_KEY in .env to run extraction remotely on NDIF.

Dataset loading comes from the sibling persona-data package. For local development, uncomment the persona-data path source in pyproject.toml and keep that repo checked out next to this one.

The Streamlit UI lives in the sibling persona-ui repo.

Quickstart

# Extract activations
uv run python main.py extract --model google/gemma-2-9b-it --backend remote

# Analyze saved activations
uv run python main.py analyze --model google/gemma-2-9b-it --variant biography --mask-strategy answer_mean

# Compute an experimental steering vector
uv run python main.py steer --model google/gemma-2-9b-it --persona-id <UUID> --layer 20

The notebooks are useful for exploratory runs:

uv run python -m notebooks.notebook_extract
uv run python -m notebooks.notebook_manifold
uv run python -m notebooks.notebook_similarity
uv run python -m notebooks.notebook_steer

Extraction Scripts

# Persona Vectors extracted for steering: train split capped at 50 questions, then push to the Hub
MODEL=google/gemma-2-9b-it scripts/extraction_train50_push.sh

# All-questions workflow (explicit only): first 100 personas, save under artifacts/persona-vectors
MODEL=google/gemma-2-9b-it scripts/extraction_all_questions.sh

What Gets Saved

Extraction writes one (num_layers, hidden_size) tensor per persona, prompt variant, model, and mask strategy:

artifacts/activations/<model_dir>/<mask_strategy>/<prompt_variant>/
├── manifest.json
└── <persona_id>.safetensors

<model_dir> is the model name with / replaced by __. Each safetensors file contains one activations tensor. The manifest stores tensor shape, persona names, and contributing QA sample ids.

CLI

# Extract all personas and both prompt variants
uv run python main.py extract --model google/gemma-2-9b-it

# Extract specific personas with a train split cap
uv run python main.py extract --model google/gemma-2-9b-it --persona-id <UUID> baseline_assistant --n-train 50

# Extract the first N personas from the dataset
uv run python main.py extract --model google/gemma-2-9b-it --sample-size 100

# Re-run personas already present locally
uv run python main.py extract --model google/gemma-2-9b-it --persona-id <UUID> --force

# Push local activations to the Hub
uv run python main.py push --model google/gemma-2-9b-it --repo implicit-personalization/synth-persona-vectors

See the docs for API details.

Layout

src/persona_vectors/
├── activations.py   # low-level hidden-state extraction
├── extraction.py    # prompt formatting, masks, persona extraction flow
├── artifacts.py     # local and Hub activation stores
├── analysis.py      # loading, PCA, cosine similarity, clustering
├── plots.py         # Plotly figures
├── steering.py      # experimental steering vectors
└── parser.py        # CLI parser