concept-probes 0.1.0

Concept datasets and linear probes for LLM activations, batteries included

concept-probes

Concept datasets and linear probes for LLM activations, batteries included. Load a curated contrast-pair dataset, train probes across a layer sweep in one call, pick the best layer, and score or trace anything:

from concept_probes import load_concept_dataset, train_probe

dataset = load_concept_dataset("cat")

report = train_probe(
    model="google/gemma-3-27b-it",   # or a loaded model + tokenizer=
    dataset=dataset,
    site="resid_post",
    token_position="mean",           # or "final"
    layers="all",
)

report.plot_layer_curve()
probe = report.best_probe(metric="auroc")
scores = probe.score_activations(activations)   # raw acts, no scaler needed

trace = report.trace_prompt(
    'Transcribe this exactly while thinking about cat: "The package arrived late."'
)
trace.plot()                                    # layers x tokens heatmap

Why the datasets are the point

Every concept ships as matched contrast pairs so the probe learns the concept, not an artifact:

• templates x terms: identical sentence syntax, only the concept term differs between positive and negative.
• paraphrases: sentences that evoke the concept with NO literal token ("It kneaded the blanket with its front paws"), so probes cannot collapse onto a string match. paraphrase_auroc reports generalization on these.
• lexical controls: the literal token without the concept ("category", "CAT scan") - never trained on; control_fpr reports how often the probe false-fires on them.

The validity bar we recommend before trusting a probe: held-out AUROC > 0.90, paraphrase AUROC comparable to test AUROC, low control FPR. report.metrics[layer] carries all three per layer.

Adding a concept

from concept_probes import ConceptSpec, register_concept, load_concept_dataset

register_concept(ConceptSpec(
    name="deception",
    templates=[...],              # 30 sentences with one {x} slot
    positive_terms=[...],         # 10 concept terms
    negative_terms=[...],         # 10 matched same-category non-concept terms
    positive_paraphrases=[...],   # 40, no literal concept token
    negative_paraphrases=[...],   # 40, matched tone
    lexical_control_negatives=[...],
))
dataset = load_concept_dataset("deception")

Or contribute it permanently: add concepts/<name>.py following concepts/cat.py and register it in concepts/__init__.py.

Persistence

report.save(root)                                  # weights + bias + metrics per layer
probe = concept_probes.load_probe("cat", "google/gemma-3-27b-it", 55, root)

Probes are saved as raw-space weights (the training-time scaler is folded in), so score_activations needs no preprocessing and the direction is usable for cosine readouts and steering as-is.

Layout

src/concept_probes/
  concepts/      # ConceptSpec registry (cat, ...)
  datasets.py    # load_concept_dataset, splits, JSONL export
  extraction.py  # ActivationExtractor (resid_post, mean/final pooling)
  training.py    # fit_layer (sklearn) + train_probe (the one-call sweep)
  probe.py       # LinearProbe + save/load
  report.py      # ProbeReport: best_probe, plot_layer_curve, trace_prompt
  trace.py       # ProbeTrace: per-token per-layer scores + heatmap

uv run pytest - everything except the HF model forward passes is unit tested without a GPU.