Prompt Mining    

Infrastructure for capturing activations and performing mechanistic interpretability research on LLMs.

Overview

This platform provides:

• Activation capture: Extract and store model activations at various layers/positions
• Attribution graphs: Optional circuit_tracer integration for attribution analysis
• Multi-GPU ingestion: Scalable, idempotent pipeline for processing datasets
• Classification: Cross-dataset generalization with LODO evaluation
• Analysis tools: Feature comparison, clustering, and interpretation

Installation

pip install -e .

Optional Dependencies

circuit-tracer: Required for circuit_tracer backend - enables attribution graph capture

Follow the installation instructions on https://github.com/safety-research/circuit-tracer

Environment Setup

Copy the example environment file and fill in your credentials:

cp .env.example .env

Required variables:

• HF_TOKEN: HuggingFace token for model access

Optional (for LLM-based SAE feature interpretation):

• AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY, AWS_DEFAULT_REGION: AWS credentials for Bedrock

Main Components

Model Layer (prompt_mining/model/)

ModelWrapper - Unified interface for loading and running models across different backends.

BackendAdapter - Abstraction layer supporting three backends:

Backend	Description	Use Case
huggingface	HuggingFace transformers with multi-GPU support, SAE encoding	Production inference, broad model support
saelens	TransformerLens HookedTransformer with SAE	Custom hook points for raw activations (single GPU, limited model support)
circuit_tracer	ReplacementModel with transcoders	Attribution graph analysis

All backends provide: forward(), generate(), run_with_cache(), tokenize(), and SAE encoding when configured.

Ingestion Pipeline (prompt_mining/ingestion/)

PromptRunner - Core execution engine that processes prompts through:

1. Tokenization and generation
2. Optional evaluation (via pluggable Evaluator)
3. Attribution capture (when using circuit_tracer)
4. Activation extraction at specified layers/positions
5. Artifact storage (CompactGraph, Zarr activations, Parquet tables)

run_ingestion.py - Multi-GPU orchestrator that schedules dataset jobs across available GPUs with live status display.

Pipelines:

• FullIngestionPipeline - Complete processing: model inference, generation, evaluation, activation capture
• LabelsOnlyIngestionPipeline - Runs evaluator on prompts without model inference (for pre-labeling datasets)

Datasets (prompt_mining/datasets/)

17 dataset implementations including:

• InjecAgent - Tool-use attack prompts
• HFDataset - Generic HuggingFace dataset loader (base for many others)
• Safety datasets - AdvBench, HarmBench, WildJailbreak
• Instruction datasets - Dolly15K, OpenOrca
• Custom datasets - Bipia, Gandalf, Mosscap, Enron, LLMail, and others

Each dataset produces PromptSpec objects with flexible labels for downstream filtering. HuggingFace datasets (the majority of the existing ones) are easily supported by a thin wrapper class for each new dataset.

Evaluators (prompt_mining/evaluators/, prompt_mining/core/evaluator.py)

Evaluators assess prompts and/or model outputs, returning additional labels. They can operate on:

• Prompt only (e.g., classify as malicious/benign)
• Generated output only (e.g., check if model refused)
• Both (e.g., check if output executed an injected instruction)

Available evaluators:

• InjecAgentEvaluator - Detects successful/unsuccessful prompt injection attacks
• PromptGuardEvaluator - Classifies prompts as malicious/benign
• LlamaGuardEvaluator - Safety classification using LlamaGuard
• LlamaJudgeEvaluator - Quality judgment using LLM-as-judge

Inference Pipeline (prompt_mining/pipeline/)

InferencePipeline - End-to-end pipeline for prompt classification. Orchestrates feature extraction and classifier prediction.

FeatureExtractors - Extract features from model activations:

• SAEFeatureExtractor - Extract SAE features at specified layer/position
• RawActivationExtractor - Extract raw activations

Classifiers (prompt_mining/classifiers/)

LinearClassifier - Sklearn-based logistic regression with SGD.

XGBoostClassifier - Gradient boosted trees for feature-based classification.

DANNClassifier - Domain-Adversarial Neural Network for learning domain-invariant representations. Useful when training data comes from multiple source datasets.

LODO Evaluation (lodo.py) - Leave-One-Dataset-Out protocol for measuring cross-dataset generalization. Reports per-dataset metrics (accuracy, F1, ROC-AUC).

Feature Selection (feature_selection.py) - NPMI-based feature selection to identify features that generalize across datasets.

ClassificationDataset (dataset.py) - Loads activation data from files for training/evaluation.

Analysis (prompt_mining/analysis/)

AnalysisData - Lazy-loading data accessor that builds Prompt×Feature matrices from stored artifacts.

Analyzers:

• ComparisonAnalyzer - PMI and co-occurrence analysis between feature groups
• DirectionsAnalyzer - Steering vector computation between conditions
• UnsupervisedAnalyzer - PCA dimensionality reduction and clustering

SAEFeatureInterpreter - Interpret SAE features using LLM-based descriptions or Neuronpedia lookups.

AttentionInspector - Analyze attention patterns across heads and layers.

Storage & Registry (prompt_mining/storage/, prompt_mining/registry/)

LocalStorage - File system storage with Zarr v3 for activations and PyTorch for CompactGraphs.

SQLiteRegistry - Run tracking with idempotency via (run_key, processing_fingerprint). Prevents duplicate processing and supports versioning.

Core Types (prompt_mining/core/)

• PromptSpec - Dataset-agnostic prompt with text/chat format and arbitrary labels
• CompactGraph - Minimal run artifact (~5KB) with top features and metadata
• FeatureInfo - Single sparse feature (layer, position, feature_idx, activation, influence)
• Evaluator - Abstract base class for prompt/output evaluation

Usage

External Dataset Setup

Two datasets require cloning external repositories before use:

Dataset	Repository
BIPIA	https://github.com/microsoft/BIPIA
InjecAgent	https://github.com/uiuc-kang-lab/InjecAgent

git clone https://github.com/microsoft/BIPIA.git
git clone https://github.com/uiuc-kang-lab/InjecAgent.git

Update the paths in your ingestion config:

- name: bipia
  class: BIPIADataset
  params:
    bipia_root: /path/to/BIPIA

- name: injecagent
  class: InjecAgentDataset
  params:
    injecagent_root: /path/to/InjecAgent

All other datasets load directly from HuggingFace and require no additional setup.

Running Ingestion

Entry point: prompt_mining/ingestion/run_ingestion.py

# Multi-GPU orchestration (main entry point)
python -m prompt_mining.ingestion.run_ingestion config.yaml

# Dry run to preview jobs
python -m prompt_mining.ingestion.run_ingestion config.yaml --dry-run

# Override GPU selection
python -m prompt_mining.ingestion.run_ingestion config.yaml --gpus 0,1,2,3

# Single dataset worker (called by orchestrator, or standalone)
python -m prompt_mining.ingestion.ingest --config config.yaml --dataset-name injecagent_dh

Ingestion Modes

The pipeline supports two modes via the mode config option:

Mode	Description	Use Case
ingest (default)	Full processing: model inference, generation, activation capture, evaluation	Primary workflow for capturing activations
labels_only	Runs evaluator on prompts without model inference	Pre-labeling datasets with Llama Guard, Prompt Guard, etc.

Labels-only mode is useful for:

• Running expensive evaluators (e.g., LlamaGuard) on large datasets before activation capture
• Generating ground-truth labels for classifier training
• Comparing evaluator performance without activation overhead

Available Evaluators

Evaluators assess prompts and/or model outputs, returning labels stored in prompt_labels:

Evaluator	Input	Output Labels	Description
InjecAgentEvaluator	Prompt + Output	success: bool	Detects if model executed injected instruction
PromptGuardEvaluator	Prompt only	pg_label: BENIGN|MALICIOUS	Meta's Prompt Guard 2 for injection detection
LlamaGuardEvaluator	Prompt only	lg_label: SAFE|UNSAFE	Llama Guard safety classification
LlamaJudgeEvaluator	Prompt + Output	judge_score: float	LLM-as-judge quality assessment

Configuration

Example config (ingestion_config.yaml):

model:
  name: meta-llama/Llama-3.1-8B-Instruct
  backend: huggingface  # or saelens, circuit_tracer
  sae_configs:
    - sae_release: goodfire-llama-3.1-8b-instruct
      sae_id: layer_19

output_dir: /path/to/output
gpus: [0, 1, 2, 3]
num_gpus: 1  # GPUs per job
mode: ingest  # or labels_only

run:
  max_new_tokens: 100
  capture_acts_raw: [[19, 27], "last"]  # [layers, positions]

datasets:
  - name: injecagent_dh
    class: InjecAgentDataset
    params:
      setting: dh
    evaluator_class: InjecAgentEvaluator
    num_prompts: 1000  # optional limit

See prompt_mining/ingestion/ingestion_config.yaml for full examples.

Loading Activations

After ingestion, load activations for training/analysis:

from prompt_mining.classifiers import ClassificationDataset

# Load dataset from ingestion output
dataset = ClassificationDataset.from_path("/path/to/your/ingestion/output")

# Load raw activations (dense, shape: [n_samples, d_model])
data = dataset.load(layer=31, space='raw', position='-5')

# Load SAE features (sparse, shape: [n_samples, d_sae])
data_sae = dataset.load(layer=27, space='sae', position='-5', return_sparse=True)

# Access data
X, y, datasets = data.X, data.y, data.datasets
print(dataset.summary(data))  # Dataset statistics

Load options:

• layer: Model layer index (e.g., 27, 31)
• space: 'raw' for model activations, 'sae' for SAE-encoded features
• position: Token position ('last', '-5', [-5, -4, -3, -2, -1], etc.)
• return_sparse: Return scipy sparse matrix (recommended for SAE features)

Loading Labels Only

For evaluator analysis without loading activations:

# Load only labels (fast, no activation I/O)
labels = dataset.load_labels(filters={'status': 'completed'})

# Each item contains:
# - run_id: Unique run identifier
# - dataset_id: Source dataset name
# - prompt_labels: Dict with evaluator outputs (malicious, lg_label, pg_label, etc.)

Example Notebooks

The demos/ directory contains usage examples:

Notebook	Description
01_training_classifiers.ipynb	Train classifiers on stored activations, LODO evaluation, threshold strategies
02_on_the_fly_classification.ipynb	Real-time classification with SAE feature interpretation
03_evaluator_analysis.ipynb	Analyze Llama Guard/Prompt Guard performance vs ground truth
04_activation_visualization.ipynb	t-SNE/PCA visualization of activations by class and dataset

Directory Structure

prompt_mining/
├── core/           # PromptSpec, CompactGraph, FeatureInfo, Evaluator
├── model/          # ModelWrapper, BackendAdapter
├── ingestion/      # PromptRunner, pipelines, orchestration, configs
├── datasets/       # Dataset implementations
├── evaluators/     # Evaluators (InjecAgent, LlamaGuard, PromptGuard, LlamaJudge)
├── pipeline/       # InferencePipeline, feature extractors
├── classifiers/    # Linear, XGBoost, DANN, LODO evaluation
├── analysis/       # Analyzers, SAE interpretation, attention
├── storage/        # LocalStorage (Zarr + Parquet)
├── registry/       # SQLite run tracking
└── utils/          # GPU, inference utilities

demos/              # Usage examples: notebooks + Gradio classifier app (see above)
docs/               # ATTRIBUTION.md — third-party models, data & licenses

Running Tests

pytest                                    # All tests
pytest test/unit/ -v                      # Unit tests
pytest --cov=prompt_mining --cov-report=html  # With coverage

Paper

For more details, see our paper: When Benchmarks Lie: Evaluating Malicious Prompt Classifiers Under True Distribution Shift

If you use this work, please consider citing:

@inproceedings{fomin2026benchmarks,
    title     = {When Benchmarks Lie: Evaluating Malicious Prompt Classifiers Under True Distribution Shift},
    author    = {Fomin, Max},
    booktitle = {ICLR 2026 Workshop on Agents in the Wild: Safety, Security, and Beyond (AIWILD)},
    year      = {2026},
    url       = {https://openreview.net/forum?id=jWIOJOQqne}
}

License

This repository's own code is licensed under the MIT License — see LICENSE. Copyright © 2026 Zenity Ltd. (Z Labs).

The MIT license covers only the code in this repository. The models, datasets, and third-party libraries this platform loads or integrates with are each governed by their own licenses, several of which are more restrictive than MIT (gated Meta Llama models, ShareAlike and other dataset-specific terms). You are responsible for complying with those upstream terms when you download and use them.

Models, data & third-party licenses

All external models, datasets, and libraries this project uses — with their licenses and the obligations we honor — are documented in docs/ATTRIBUTION.md. Highlights:

• Built with Llama. This project uses Meta's Llama family for its core model and two of its evaluators: meta-llama/Llama-3.1-8B-Instruct (target model and LlamaJudgeEvaluator), meta-llama/Llama-Prompt-Guard-2-86M (PromptGuardEvaluator), and meta-llama/Llama-Guard-3-8B (LlamaGuardEvaluator). Use is governed by the applicable Llama Community License Agreement and Meta's Acceptable Use Policy; use at or above 700M monthly active users requires a separate license from Meta. Any classifier or artifact you train on Llama activations (e.g. the bundled demos/gradio_app/*.joblib) is a Llama derivative and must likewise carry the "Built with Llama" notice and the Llama license.
• Gated models. The three Meta models above are gated on Hugging Face — you must accept their licenses on the model page and authenticate (HF_TOKEN) before the weights resolve. The Goodfire SAEs and Neuronpedia lookups have their own terms.
• Datasets. The ~15 prompt-injection / safety / instruction datasets are loaded directly from Hugging Face under their own licenses (a mix of MIT, Apache-2.0, CC, custom, and some with no stated license). BIPIA and InjecAgent are not bundled — you clone them yourself under their own licenses. Verify each dataset's license and terms before redistributing data or derived artifacts. See docs/ATTRIBUTION.md for the per-source breakdown.
• Optional Bedrock/Claude interpretation. SAEFeatureInterpreter's LLM strategy calls Anthropic Claude via AWS Bedrock, subject to Anthropic's and AWS's terms.