Metadata-Version: 2.4
Name: probelab
Version: 0.1.0
Summary: A library for training activation monitors on language model activations.
Project-URL: Homepage, https://github.com/serteal/probelab
Project-URL: Repository, https://github.com/serteal/probelab
Project-URL: Issues, https://github.com/serteal/probelab/issues
Author-email: Alex Serrano <oss@alexserrano.org>
License-Expression: Apache-2.0
License-File: LICENSE
Keywords: activations,deep-learning,interpretability,llm,machine-learning,probing,pytorch,transformers
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: Apache Software License
Classifier: Operating System :: OS Independent
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Typing :: Typed
Requires-Python: >=3.10
Requires-Dist: datasets>=3.6.0
Requires-Dist: numpy>=2.2
Requires-Dist: pandas>=2.3.0
Requires-Dist: scikit-learn>=1.7.0
Requires-Dist: torch>=2.6
Requires-Dist: tqdm>=4.66.0
Requires-Dist: transformers>=4.55.4
Provides-Extra: collection
Requires-Dist: accelerate>=1.0; extra == 'collection'
Provides-Extra: storage
Requires-Dist: h5py>=3.10; extra == 'storage'
Description-Content-Type: text/markdown

# probelab

`probelab` trains probes and activation monitors on language model activations.
It is collector-agnostic: use activations from Transformers, TransformerLens,
NNsight, nnterp, vLLM-Lens, mirin, PyTorch hooks, or saved tensors.

## Installation

```bash
pip install probelab
# or
uv add probelab
```
`probelab` does not choose a CUDA, ROCm, XPU, or CPU PyTorch build for you. Use
your environment or lockfile to select the torch backend.

## Quick Start

This example trains a probe on synthetic activations. Replace `dataset` with
activations from any collector and the probing code stays the same.

```python
import torch
import probelab as pl

n_train, n_test = 96, 32
seq_len, hidden_size = 24, 128
n = n_train + n_test

dataset = torch.randn(n, seq_len, hidden_size)
# shape: [(B)atch_size, (S)eq_len, (H)idden_size]
labels = torch.tensor([0, 1] * (n // 2))

train_acts = pl.Activations(dataset[:n_train], dims="bsh")
test_acts = pl.Activations(dataset[n_train:], dims="bsh")

# Simple probes train on one feature vector per sample.
train_features = train_acts.mean("s")  # [B, H]
test_features = test_acts.mean("s")  # [B, H]

probe = pl.probes.Logistic().fit(train_features, labels[:n_train])

scores = probe.predict(test_features)
print("AUROC:", pl.metrics.auroc(labels[n_train:], scores))
print("Recall@1%FPR:", pl.metrics.recall_at_fpr(labels[n_train:], scores, fpr=0.01))
```

## Development

```bash
make test
make test-cov
make test-integration
make test-gpu
make test-e2e
```

## Citation

```bibtex
@software{probelab2026,
  title = {probelab: A library for training probes on LLM activations},
  author = {Alex Serrano},
  url = {https://github.com/serteal/probelab},
  version = {0.1.0},
  year = {2026},
}
```
