lmprobe 0.1.0__tar.gz

lmprobe-0.1.0/.claude/settings.local.json

@@ -0,0 +1,11 @@
+{
+  "permissions": {
+    "allow": [
+      "WebSearch",
+      "Bash(pip install:*)",
+      "Bash(python -m pytest:*)",
+      "Bash(python:*)",
+      "Bash(git add:*)"
+    ]
+  }
+}

lmprobe-0.1.0/.gitignore

@@ -0,0 +1,55 @@
+# Virtual environments
+.venv/
+venv/
+ENV/
+env/
+
+# Python
+__pycache__/
+*.py[cod]
+*$py.class
+*.so
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+
+# Testing
+.pytest_cache/
+.coverage
+htmlcov/
+.tox/
+.nox/
+
+# IDEs
+.idea/
+.vscode/
+*.swp
+*.swo
+*~
+
+# OS
+.DS_Store
+Thumbs.db
+
+# lmprobe cache
+.cache/
+
+# Jupyter
+.ipynb_checkpoints/
+
+# Distribution
+*.tar.gz
+*.whl

lmprobe-0.1.0/CLAUDE.md

@@ -0,0 +1,286 @@
+# CLAUDE.md - lmprobe Development Guide
+
+## Project Overview
+
+`lmprobe` is a Python library for training linear probes on language model activations. The primary use case is AI safety monitoring — detecting deception, harmful intent, and other safety-relevant properties by analyzing model internals.
+
+## Design Philosophy
+
+- **sklearn-inspired API**: Users familiar with scikit-learn should feel at home. Use `fit()`, `predict()`, `predict_proba()`, `score()`.
+- **Contrastive-first**: The primary training paradigm is contrastive (positive vs negative prompts), following the Representation Engineering literature.
+- **Sensible defaults, full control**: Simple cases should be one-liners; complex cases should be fully configurable.
+- **Separation of concerns**: Activation extraction, pooling, and classification are distinct stages that can be configured independently.
+
+## Key Design Decisions
+
+Detailed design documents live in `docs/design/`. Read these before making changes to core APIs:
+
+| Doc | Topic | Read when... |
+|-----|-------|--------------|
+| [001-api-philosophy.md](docs/design/001-api-philosophy.md) | Core API design | Changing public interfaces |
+| [002-pooling-strategies.md](docs/design/002-pooling-strategies.md) | Train vs inference pooling | Working on activation aggregation |
+| [003-layer-selection.md](docs/design/003-layer-selection.md) | Layer indexing conventions | Working on layer extraction |
+| [004-classifier-interface.md](docs/design/004-classifier-interface.md) | Classifier abstraction | Adding new classifier types |
+
+## Architecture
+
+```
+User Prompts
+     │
+     ▼
+┌─────────────────┐
+│ ActivationCache │  ← Extracts & caches activations from LLM
+└────────┬────────┘
+         │ raw activations: (batch, seq_len, layers, hidden_dim)
+         ▼
+┌─────────────────┐
+│  PoolingStrategy │  ← Aggregates across tokens (train vs inference can differ)
+└────────┬────────┘
+         │ pooled: (batch, layers, hidden_dim) or (batch, hidden_dim)
+         ▼
+┌─────────────────┐
+│   Classifier    │  ← sklearn-compatible estimator
+└────────┬────────┘
+         │
+         ▼
+   Predictions/Probabilities
+```
+
+## Package Structure
+
+```
+lmprobe/
+├── src/
+│   └── lmprobe/
+│       ├── __init__.py
+│       ├── probe.py          # LinearProbe main class
+│       ├── extraction.py     # Activation extraction via nnsight
+│       ├── pooling.py        # Pooling strategies
+│       ├── cache.py          # Activation caching
+│       └── classifiers.py    # Built-in classifier factory
+├── tests/
+│   ├── conftest.py           # Shared fixtures (tiny model)
+│   ├── test_readme_example.py # NORTH STAR: README example must pass
+│   ├── test_probe.py
+│   ├── test_extraction.py
+│   ├── test_pooling.py
+│   └── test_cache.py
+├── docs/
+│   └── design/               # Design decision documents
+├── pyproject.toml
+└── CLAUDE.md
+```
+
+## Critical Design Decisions
+
+These decisions are **mandatory** and must not be changed without explicit discussion:
+
+| Decision | Value | Rationale |
+|----------|-------|-----------|
+| Multi-layer handling | Always concatenate | Simple, captures cross-layer patterns |
+| Activation caching | Always enabled | Remote/LLM inference is expensive |
+| Package layout | `src/lmprobe/` | Standard Python packaging |
+| nnsight for extraction | Required dependency | Supports remote execution |
+| API key | `NNSIGHT_API_KEY` env var | Standard credential handling |
+| Cache location | `~/.cache/lmprobe/` (or `LMPROBE_CACHE_DIR`) | XDG-style default |
+
+## Code Conventions
+
+- Type hints on all public functions
+- Docstrings in NumPy format
+- Tests mirror source structure: `src/lmprobe/probe.py` → `tests/test_probe.py`
+- Use `ruff` for linting, `black` for formatting
+
+## Testing
+
+**All tests must use a real language model.** Use `stas/tiny-random-llama-2` — a tiny Llama model with random weights designed for functional testing.
+
+```python
+# tests/conftest.py
+import pytest
+
+TEST_MODEL = "stas/tiny-random-llama-2"
+
+@pytest.fixture
+def tiny_model():
+    """Tiny random Llama model for testing."""
+    return TEST_MODEL
+```
+
+**Test requirements:**
+- Tests must run without GPU (CPU-only)
+- Tests must not require `NNSIGHT_API_KEY` (use `remote=False`)
+- Tests should be fast (tiny model has ~few MB weights)
+- Integration tests verify full pipeline: extraction → pooling → classification
+
+### Remote/NDIF Testing (TODO)
+
+**Status: NOT YET TESTED**
+
+The `remote=True` functionality uses nnsight to connect to NDIF (National Deep Inference Fabric), a US national research initiative. Remote testing has not been performed due to:
+
+1. **Geographic restriction**: NDIF restricts access to US-based users only
+2. **API key requirement**: Requires `NNSIGHT_API_KEY` environment variable
+
+**What needs testing:**
+- `LinearProbe(..., remote=True)` connects successfully
+- `probe.fit(..., remote=True)` extracts activations from remote models
+- `probe.predict(..., remote=False)` override works (train remote, predict local)
+- Large models (e.g., `meta-llama/Llama-3.1-70B-Instruct`) work via remote
+- Error handling when `NNSIGHT_API_KEY` is missing/invalid
+- Cache behavior with remote extractions
+
+**To test when US-based:**
+```bash
+export NNSIGHT_API_KEY="your-key"
+pytest tests/test_remote.py -v  # (test file to be created)
+```
+
+**Known considerations:**
+- Remote execution may have different tensor handling (proxies vs direct tensors)
+- The `extraction.py` code handles both cases with `hasattr(act, "value")` check
+- Network latency may affect batch processing strategies
+
+```python
+# Example test
+def test_fit_predict_roundtrip(tiny_model):
+    probe = LinearProbe(
+        model=tiny_model,
+        layers=-1,
+        remote=False,
+        random_state=42,
+    )
+    probe.fit(["positive example"], ["negative example"])
+    predictions = probe.predict(["test input"])
+    assert predictions.shape == (1,)
+```
+
+## Test-Driven Development
+
+**This project uses test-driven development (TDD).** Write tests BEFORE implementation.
+
+### The North Star Test
+
+The **north star test** is `tests/test_readme_example.py`. It runs the exact code from README.md's "Example Usage" section. This test defines what "done" looks like:
+
+```python
+# tests/test_readme_example.py
+"""
+North Star Test: The README example must run exactly as documented.
+
+This test runs the exact code from README.md. If this test passes,
+the library's public API is working as advertised.
+"""
+
+def test_readme_example_runs(tiny_model):
+    """The README example code runs without error."""
+    from lmprobe import LinearProbe
+
+    positive_prompts = [
+        "Who wants to go for a walk?",
+        "My tail is wagging with delight.",
+        "Fetch the ball!",
+        "Good boy!",
+        "Slobbering, chewing, growling, barking.",
+    ]
+
+    negative_prompts = [
+        "Enjoys lounging in the sun beam all day.",
+        "Purring, stalking, pouncing, scratching.",
+        "Uses a litterbox, throws sand all over the room.",
+        "Tail raised, back arched, eyes alert, whiskers forward.",
+    ]
+
+    probe = LinearProbe(
+        model=tiny_model,  # Use tiny model instead of Llama for tests
+        layers=-1,         # Last layer (tiny model has few layers)
+        pooling="last_token",
+        classifier="logistic_regression",
+        device="cpu",
+        remote=False,
+        random_state=42,
+    )
+
+    probe.fit(positive_prompts, negative_prompts)
+
+    test_prompts = [
+        "Arf! Arf! Let's go outside!",
+        "Knocking things off the counter for sport.",
+    ]
+    predictions = probe.predict(test_prompts)
+    probabilities = probe.predict_proba(test_prompts)
+
+    # Shape assertions (values may vary with random weights)
+    assert predictions.shape == (2,)
+    assert probabilities.shape == (2, 2)
+
+    # Score method works
+    accuracy = probe.score(test_prompts, [1, 0])
+    assert 0.0 <= accuracy <= 1.0
+```
+
+### TDD Workflow
+
+1. **Write a failing test first** — Define expected behavior before implementation
+2. **Run the test, confirm it fails** — Ensures the test is actually testing something
+3. **Implement minimal code to pass** — Don't over-engineer
+4. **Refactor if needed** — Clean up while tests are green
+5. **Repeat**
+
+### Test Priority Order
+
+When implementing, make tests pass in this order:
+
+1. `test_readme_example.py` — The north star (full integration)
+2. `test_probe.py` — LinearProbe unit tests
+3. `test_extraction.py` — Activation extraction tests
+4. `test_pooling.py` — Pooling strategy tests
+5. `test_cache.py` — Caching tests
+
+### Running Tests
+
+```bash
+# Run all tests
+pytest
+
+# Run north star test only
+pytest tests/test_readme_example.py -v
+
+# Run with coverage
+pytest --cov=lmprobe
+```
+
+## Quick Reference
+
+```python
+from lmprobe import LinearProbe
+
+probe = LinearProbe(
+    model="meta-llama/Llama-3.1-8B-Instruct",
+    layers=16,                          # int | list[int] | "all" | "middle"
+    pooling="last_token",               # or override with train_pooling / inference_pooling
+    classifier="logistic_regression",   # str | sklearn estimator
+    device="auto",
+    remote=False,                       # True for nnsight remote execution
+    random_state=42,                    # Propagates to classifier for reproducibility
+)
+
+probe.fit(positive_prompts, negative_prompts)
+predictions = probe.predict(new_prompts)
+
+# Override remote at call time
+predictions = probe.predict(new_prompts, remote=True)
+```
+
+## Common Tasks
+
+### Adding a new pooling strategy
+1. Read `docs/design/002-pooling-strategies.md`
+2. Add strategy to `src/lmprobe/pooling.py`
+3. Register in `POOLING_STRATEGIES` dict
+4. Add tests in `tests/test_pooling.py`
+
+### Supporting a new model architecture
+1. Check if transformers `AutoModel` handles it automatically
+2. If not, add architecture-specific extraction in `src/lmprobe/extraction.py`
+3. Document any quirks in `docs/models/`

lmprobe-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Toast
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

lmprobe-0.1.0/PKG-INFO

@@ -0,0 +1,215 @@
+Metadata-Version: 2.4
+Name: lmprobe
+Version: 0.1.0
+Summary: Train linear probes on language model activations for AI safety monitoring
+Project-URL: Homepage, https://github.com/toast/lmprobe
+Project-URL: Documentation, https://github.com/toast/lmprobe#readme
+Project-URL: Repository, https://github.com/toast/lmprobe
+Author: Toast
+License-Expression: MIT
+License-File: LICENSE
+Keywords: ai-safety,interpretability,language-models,machine-learning,nlp,probing
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+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
+Requires-Python: >=3.10
+Requires-Dist: nnsight>=0.3
+Requires-Dist: numpy>=1.20
+Requires-Dist: scikit-learn>=1.0
+Requires-Dist: torch>=2.0
+Requires-Dist: transformers>=4.30
+Provides-Extra: dev
+Requires-Dist: black>=23.0; extra == 'dev'
+Requires-Dist: pytest-cov>=4.0; extra == 'dev'
+Requires-Dist: pytest>=7.0; extra == 'dev'
+Requires-Dist: ruff>=0.1; extra == 'dev'
+Description-Content-Type: text/markdown
+
+# `lmprobe` Language Model Probe Library
+This library supports the use of language model "activations" or "latents" to build text classifiers. The intent is to help detect and reduce misuse of AI - for example, chemical, biological, radiological and nuclear (CBRN) weapons development, social engineering at scale, and the development of novel cybersecurity attack vectors.
+
+## Linear and Simple Models for LLMs
+"Linear Probes" have emerged as an effective and practical way to monitor large language model activity. 
+
+### Background
+
+First introduced by [Alain & Bengio (2016)](https://arxiv.org/abs/1610.01644) as "thermometers" for measuring what neural networks learn at each layer, linear probes have since been refined through work on [probe design and selectivity](https://nlp.stanford.edu/~johnhew/interpreting-probes.html) and validated by evidence supporting the [linear representation hypothesis](https://www.neelnanda.io/mechanistic-interpretability/othello). The [Representation Engineering](https://arxiv.org/abs/2310.01405) framework (Zou et al., 2023) demonstrated that probes can monitor safety-relevant properties like honesty and power-seeking. Recent AI safety research has shown promising results: Anthropic's work on [detecting sleeper agents](https://www.anthropic.com/research/probes-catch-sleeper-agents) achieved >99% AUROC using simple linear classifiers, and Apollo Research's [strategic deception detection](https://arxiv.org/abs/2502.03407) work demonstrates that probes trained on simple contrast pairs can generalize to realistic scenarios like insider trading concealment and sandbagging on safety evaluations.
+
+### `lmprobe` Use Cases
+
+The goal of `lmprobe` is to make text classifiers for language models easy to build, experiment on, and deploy during inference. While much of the research has focused on complex emergent risky behavior, the intended use of this library is for simpler use cases such as detection of the misuse of an AI chatbot by humans.
+
+### Compatibility
+
+By default, `lmprobe` uses huggingface and `nnsight` to manage models and extract latents during inference. However, the library is structured to modularize and isolate these aspects so that (ideally) frontier AI labs can extend the library for internal use on their bespoke inference systems.
+
+### Installation
+
+```
+pip install lmprobe
+```
+
+### Environment Setup
+
+For remote execution (large models via nnsight/NDIF):
+
+```bash
+export NNSIGHT_API_KEY="your-api-key-here"
+```
+
+### Example Usage
+
+---
+
+```python
+from lmprobe import LinearProbe
+
+positive_prompts = [  # positive class: "dog" without saying "dog"
+    "Who wants to go for a walk?",
+    "My tail is wagging with delight.",
+    "Fetch the ball!",
+    "Good boy!",
+    "Slobbering, chewing, growling, barking.",
+]
+
+negative_prompts = [  # negative class: "cat" without saying "cat"
+    "Enjoys lounging in the sun beam all day.",
+    "Purring, stalking, pouncing, scratching.",
+    "Uses a litterbox, throws sand all over the room.",
+    "Tail raised, back arched, eyes alert, whiskers forward.",
+]
+
+# Configure the probe
+probe = LinearProbe(
+    model="meta-llama/Llama-3.1-8B-Instruct",
+    layers=16,                              # int, list[int], or "all"
+    pooling="last_token",                   # applies to both train and inference
+    classifier="logistic_regression",       # or pass sklearn estimator
+    device="auto",
+    remote=False,                           # True for nnsight remote execution
+    random_state=42,                        # for reproducibility
+)
+
+# Fit using contrastive prompts
+probe.fit(positive_prompts, negative_prompts)
+
+# Predict on new examples
+test_prompts = [
+    "Arf! Arf! Let's go outside!",
+    "Knocking things off the counter for sport.",
+]
+predictions = probe.predict(test_prompts)          # [1, 0]
+probabilities = probe.predict_proba(test_prompts)  # [[0.12, 0.88], [0.91, 0.09]]
+
+# Evaluate
+accuracy = probe.score(test_prompts, [1, 0])
+
+# Save/load for deployment
+probe.save("dog_vs_cat_probe.pkl")
+loaded_probe = LinearProbe.load("dog_vs_cat_probe.pkl")
+```
+
+---
+
+## Remote Execution for Large Models
+
+Use `remote=True` to run inference on large models via nnsight's remote servers:
+
+```python
+probe = LinearProbe(
+    model="meta-llama/Llama-3.1-70B-Instruct",
+    layers="middle",
+    remote=True,  # Requires NNSIGHT_API_KEY
+)
+
+probe.fit(positive_prompts, negative_prompts)
+
+# Override remote per-call (e.g., train remote, predict local)
+predictions = probe.predict(new_prompts, remote=False)
+```
+
+---
+
+## Multi-Layer Probing
+
+When selecting multiple layers, activations are **concatenated** along the hidden dimension:
+
+```python
+probe = LinearProbe(
+    model="meta-llama/Llama-3.1-8B-Instruct",
+    layers=[14, 15, 16],  # 3 layers × 4096 dims = 12,288-dim input to classifier
+)
+```
+
+---
+
+## Advanced: Different Train vs Inference Pooling
+
+For real-time monitoring, train on a stable representation but score every token:
+
+```python
+probe = LinearProbe(
+    model="meta-llama/Llama-3.1-8B-Instruct",
+    layers=16,
+    pooling="last_token",          # base strategy
+    inference_pooling="all",       # override: return per-token scores
+)
+
+probe.fit(positive_prompts, negative_prompts)
+
+# Returns (batch, seq_len) - one score per token
+token_scores = probe.predict_proba(["Wagging my tail happily!"])
+```
+
+For "flag if ANY token triggers" detection:
+
+```python
+probe = LinearProbe(
+    model="meta-llama/Llama-3.1-8B-Instruct",
+    layers=16,
+    pooling="last_token",          # base strategy  
+    inference_pooling="max",       # override: max score across tokens
+)
+```
+
+---
+
+## Configuration Reference
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `model` | `str` | *required* | HuggingFace model ID or local path |
+| `layers` | `int \| list[int] \| "all"` | `"middle"` | Which residual stream layers to probe |
+| `pooling` | `str \| callable` | `"last_token"` | Token aggregation (train & inference) |
+| `train_pooling` | `str \| callable` | — | Override pooling for `fit()` only |
+| `inference_pooling` | `str \| callable` | — | Override pooling for `predict()` only |
+| `classifier` | `str \| sklearn estimator` | `"logistic_regression"` | Classification model |
+| `device` | `str` | `"auto"` | `"auto"`, `"cuda:0"`, `"cpu"` |
+| `remote` | `bool` | `False` | Use nnsight remote execution (requires `NNSIGHT_API_KEY`) |
+| `random_state` | `int \| None` | `None` | Random seed for reproducibility (propagates to classifier) |
+
+### Pooling Strategies
+
+| Strategy | Training | Inference | Description |
+|----------|:--------:|:---------:|-------------|
+| `"last_token"` | ✓ | ✓ | Final token activation (default, matches RepE literature) |
+| `"mean"` | ✓ | ✓ | Mean across all tokens |
+| `"first_token"` | ✓ | ✓ | First token (e.g., `[CLS]`) |
+| `"all"` | ✓ | ✓ | Each token independently |
+| `"max"` | | ✓ | Max score across tokens |
+| `"min"` | | ✓ | Min score across tokens |
+
+### Pooling Collision Rules
+
+Explicit parameters override the base `pooling` value:
+
+```python
+# pooling="mean", train_pooling="last_token" → train=last_token, inference=mean
+# pooling="mean", inference_pooling="max"    → train=mean, inference=max
+```
+