Metadata-Version: 2.4
Name: concept-probes
Version: 0.1.0
Summary: Concept datasets and linear probes for LLM activations, batteries included
Keywords: interpretability,linear-probes,llm,activations,concepts
Author: Joshua Fonseca
Author-email: Joshua Fonseca <jerseyfonseca@gmail.com>
License-Expression: MIT
License-File: LICENSE
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Science/Research
Classifier: Programming Language :: Python :: 3.12
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Dist: matplotlib>=3.10.0
Requires-Dist: numpy>=2.0.0
Requires-Dist: scikit-learn>=1.5.0
Requires-Dist: torch>=2.4.0
Requires-Dist: transformers>=4.44.0
Requires-Python: >=3.12
Project-URL: Homepage, https://github.com/vuciv/concept-probes
Project-URL: Repository, https://github.com/vuciv/concept-probes
Project-URL: Issues, https://github.com/vuciv/concept-probes/issues
Description-Content-Type: text/markdown

# 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:

```python
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

```python
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

```python
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.
