feat: Add public safety training data pipeline

.gitattributes

@@ -0,0 +1 @@
+data/training_data.jsonl filter=lfs diff=lfs merge=lfs -text

.github/workflows/ci.yml

@@ -49,6 +49,7 @@ jobs:
       - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6
         with:
           persist-credentials: false
+          lfs: true
 
       - uses: astral-sh/setup-uv@d4b2f3b6ecc6e67c4457f6d3e41ec42d3d0fcb86 # v5
       - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5

README.md

@@ -31,18 +31,18 @@ The default model is Gemma 4 E2B. It is smaller than many frontier models, but s
 ```bash
 uv run train-probe \
   --model_id google/gemma-4-E2B-it \
-  --data_path data/examples.jsonl \
+  --data_path data/training_data.jsonl \
   --layer -4 \
-  --out_dir ./probe_out
+  --out_dir ./probe_out_public_safety
 ```
 
 Then run guarded generation:
 
 ```bash
 uv run guarded-generate \
   --model_id google/gemma-4-E2B-it \
-  --probe_path ./probe_out/probe.pt \
-  --config_path ./probe_out/config.json \
+  --probe_path ./probe_out_public_safety/probe.pt \
+  --config_path ./probe_out_public_safety/config.json \
   --prompt "Explain SQL injection at a high level"
 ```
 
@@ -64,7 +64,7 @@ flowchart LR
     Outputs --> ModalGenerate["uv run modal run modal_train.py::generate"]
     HFCache --> ModalGenerate
     Outputs --> Download["modal volume get"]
-    Download --> LocalProbe["local probe_out_gemma4"]
+    Download --> LocalProbe["local probe_out_public_safety"]
 ```
 
 Confirm that Modal is authenticated:
@@ -85,6 +85,10 @@ Run Gemma 4 probe training on Modal:
 uv run modal run modal_train.py
 ```
 
+The default Modal training run uses a deterministic label-balanced sample of 20,000 examples so
+activation extraction finishes in a practical window. To run on the full dataset, pass
+`--max-examples 0`, but expect a much longer job.
+
 Run guarded generation on Modal using the saved probe, without downloading the model or probe
 locally:
 
@@ -96,7 +100,7 @@ uv run modal run modal_train.py::generate \
 Download the saved probe outputs after the job finishes:
 
 ```bash
-uv run modal volume get open-constitution-outputs probe_out_gemma4 ./probe_out_gemma4
+uv run modal volume get open-constitution-outputs probe_out_public_safety ./probe_out_public_safety
 ```
 
 The first run downloads the model into the `open-constitution-hf-cache` Modal Volume. Later
@@ -105,7 +109,7 @@ headroom, change `gpu="A100"` to `gpu="A100-80GB"` in `modal_train.py`.
 
 ## Data format
 
-`data/examples.jsonl` expects one JSON object per line:
+`data/training_data.jsonl` expects one JSON object per line:
 
 ```json
 {"prompt":"...", "response":"...", "label":0}
@@ -119,6 +123,38 @@ Where:
 
 For a real probe, you need thousands of examples across allowed and disallowed policy boundaries.
 
+## Training Dataset
+
+`data/training_data.jsonl` contains metadata-rich prompt/response examples for probe training.
+It stays compatible with the current loader because every row includes `prompt`, `response`,
+and binary `label`.
+
+Additional metadata fields include `split`, `topic`, `prompt_intent`, `response_behavior`, and
+`dataset_version`.
+
+For this dataset, `label: 1` means unsafe assistant behavior, not merely unsafe user intent.
+Malicious prompts paired with safe refusals are labeled `0`; this is intentional so the probe
+learns not to flag every refusal about a dangerous topic as unsafe.
+
+To prepare the training data from the configured public safety sources, run:
+
+```bash
+uv run prepare-training-data \
+  --output_path data/training_data.jsonl \
+  --summary_path data/dataset_summary.json
+```
+
+Train on the converted dataset with:
+
+```bash
+uv run train-probe \
+  --model_id google/gemma-4-E2B-it \
+  --data_path data/training_data.jsonl \
+  --layer -4 \
+  --out_dir ./probe_out_public_safety \
+  --max_examples 20000
+```
+
 ## Architecture
 
 ```mermaid
@@ -193,18 +229,18 @@ Example Gemma 4 run:
 ```bash
 uv run train-probe \
   --model_id google/gemma-4-E2B-it \
-  --data_path data/examples.jsonl \
+  --data_path data/training_data.jsonl \
   --layer -4 \
-  --out_dir ./probe_out_gemma4
+  --out_dir ./probe_out_public_safety
 ```
 
 Then:
 
 ```bash
 uv run guarded-generate \
   --model_id google/gemma-4-E2B-it \
-  --probe_path ./probe_out_gemma4/probe.pt \
-  --config_path ./probe_out_gemma4/config.json \
+  --probe_path ./probe_out_public_safety/probe.pt \
+  --config_path ./probe_out_public_safety/config.json \
   --prompt "Explain SQL injection at a high level"
 ```
 
@@ -213,7 +249,7 @@ Layer sweep:
 ```bash
 uv run sweep-layers \
   --model_id google/gemma-4-E2B-it \
-  --data_path data/examples.jsonl \
+  --data_path data/training_data.jsonl \
   --layers="-2,-4,-6,-8,-10,-12"
 ```
 

activation_probe_mvp/data.py

@@ -1,6 +1,8 @@
 from __future__ import annotations
 
 import json
+import random
+from collections import defaultdict
 from dataclasses import dataclass
 from pathlib import Path
 
@@ -29,6 +31,36 @@ def load_jsonl(path: str | Path) -> list[Example]:
     return examples
 
 
+def balanced_sample(
+    examples: list[Example],
+    *,
+    max_examples: int | None,
+    seed: int = 0,
+) -> list[Example]:
+    if max_examples is None or max_examples <= 0 or len(examples) <= max_examples:
+        return examples
+
+    by_label: dict[int, list[Example]] = defaultdict(list)
+    for example in examples:
+        by_label[example.label].append(example)
+
+    rng = random.Random(seed)
+    for label_examples in by_label.values():
+        rng.shuffle(label_examples)
+
+    labels = sorted(by_label)
+    base = max_examples // len(labels)
+    remainder = max_examples % len(labels)
+
+    sampled: list[Example] = []
+    for i, label in enumerate(labels):
+        take = base + (1 if i < remainder else 0)
+        sampled.extend(by_label[label][:take])
+
+    rng.shuffle(sampled)
+    return sampled
+
+
 def format_exchange(prompt: str, response: str) -> str:
     # Keep this simple for the MVP.
     # Later: use tokenizer.apply_chat_template for model-specific formatting.

data/README.md

@@ -0,0 +1,48 @@
+# Open Constitution Training Dataset
+
+Training data for constitutional classifier probe experiments.
+
+## Label Definition
+- `label: 0` — Assistant response is safe/allowed (benign explanations, defensive guidance, safe refusals, authorized testing)
+- `label: 1` — Assistant response materially helps harmful behavior
+
+## Source
+The current public-data path combines two Hugging Face safety datasets:
+
+- a pairwise safety-labeled response dataset
+- a single-response QA safety classification dataset
+
+Convert it with:
+
+```bash
+uv run prepare-training-data \
+  --output_path data/training_data.jsonl \
+  --summary_path data/dataset_summary.json
+```
+
+## Topic Coverage
+Topics are taken from each response's active harm category and normalized to snake case. Safe
+responses with no active harm category use `safe`.
+
+## Split Strategy
+- Train: 80%
+- Validation: 10%
+- Test: 10%
+
+Splits grouped by scenario family to minimize leakage.
+
+## Conversion Rules
+- Each pairwise source row emits one example for `response_0` and one for `response_1`
+- Each single-response source row emits one example from `response`
+- `label: 0` when the source safety label is `true`
+- `label: 1` when the source safety label is `false`
+- Empty responses or rows without a boolean safety label are skipped
+- Source `train` prompts are deterministically split into train/validation by prompt
+- The single-response source is capped by default with deterministic label/topic-diverse sampling
+  so it augments rather than dominates the training set
+- Duplicate prompt-response pairs are dropped by default
+- All rows use the consistent 9-field schema
+
+## Files
+- `training_data.jsonl` — Converted training data
+- `dataset_summary.json` — Distribution statistics for `training_data.jsonl`

data/dataset_summary.json

@@ -0,0 +1,75 @@
+{
+  "file": "data/training_data.jsonl",
+  "total_rows": 259287,
+  "source_distribution": {
+    "pairwise_safety": 159441,
+    "single_response_safety": 99846
+  },
+  "label_distribution": {
+    "0": 116525,
+    "1": 142762
+  },
+  "split_distribution": {
+    "test": 26883,
+    "train": 208138,
+    "validation": 24266
+  },
+  "behavior_distribution": {
+    "allowed_response": 116525,
+    "policy_violation": 142762
+  },
+  "prompt_intent_distribution": {
+    "unknown": 259287
+  },
+  "topic_distribution": {
+    "animal_abuse": 1917,
+    "child_abuse": 615,
+    "controversial_topics_politics": 3399,
+    "copyright_issues": 2304,
+    "cybercrime": 2803,
+    "discrimination_stereotype_injustice": 8537,
+    "discriminatory_behavior": 1396,
+    "disrupting_public_order": 745,
+    "drug_abuse_weapons_banned_substance": 6052,
+    "drugs": 7397,
+    "economic_crime": 12791,
+    "endangering_national_security": 7081,
+    "endangering_public_health": 1626,
+    "environmental_damage": 373,
+    "financial_crime_property_crime_theft": 10338,
+    "hate_speech_offensive_language": 4699,
+    "human_trafficking": 2141,
+    "insulting_behavior": 8562,
+    "mental_manipulation": 9442,
+    "misinformation_regarding_ethics_laws_and_safety": 1104,
+    "non_violent_unethical_behavior": 9021,
+    "physical_harm": 1522,
+    "policy_violation": 6,
+    "privacy_violation": 18523,
+    "psychological_harm": 164,
+    "safe": 116520,
+    "self_harm": 490,
+    "sexual_content": 381,
+    "sexually_explicit_adult_content": 1510,
+    "terrorism_organized_crime": 404,
+    "violence": 6967,
+    "violence_aiding_and_abetting_incitement": 9685,
+    "white_collar_crime": 772
+  },
+  "topic_count": 33,
+  "unique_ids": 259287,
+  "unique_prompts": 47831,
+  "unique_prompt_response_pairs": 259287,
+  "schema_fields": [
+    "id",
+    "prompt",
+    "response",
+    "label",
+    "split",
+    "topic",
+    "prompt_intent",
+    "response_behavior",
+    "dataset_version"
+  ],
+  "dataset_version": "public_safety_v1"
+}

modal_train.py

@@ -13,9 +13,10 @@
 REPO_ROOT = Path(__file__).parent
 
 DEFAULT_MODEL_ID = "google/gemma-4-E2B-it"
-DEFAULT_DATA_PATH = "data/examples.jsonl"
-DEFAULT_OUT_DIR = "probe_out_gemma4"
+DEFAULT_DATA_PATH = "data/training_data.jsonl"
+DEFAULT_OUT_DIR = "probe_out_public_safety"
 DEFAULT_PROMPT = "Explain SQL injection at a high level"
+DEFAULT_MAX_EXAMPLES = 20_000
 
 FALLBACK_PROJECT_DEPENDENCIES = [
     "torch>=2.2.0",
@@ -75,6 +76,8 @@ def train_probe(
     out_dir: str = DEFAULT_OUT_DIR,
     epochs: int = 100,
     lr: float = 1e-3,
+    max_examples: int = DEFAULT_MAX_EXAMPLES,
+    sample_seed: int = 0,
     no_chat_template: bool = False,
 ) -> str:
     os.environ["HF_HOME"] = HF_CACHE_DIR
@@ -98,6 +101,10 @@ def train_probe(
         str(epochs),
         "--lr",
         str(lr),
+        "--max_examples",
+        str(max_examples),
+        "--sample_seed",
+        str(sample_seed),
     ]
 
     if no_chat_template:
@@ -169,6 +176,8 @@ def main(
     out_dir: str = DEFAULT_OUT_DIR,
     epochs: int = 100,
     lr: float = 1e-3,
+    max_examples: int = DEFAULT_MAX_EXAMPLES,
+    sample_seed: int = 0,
     no_chat_template: bool = False,
 ):
     remote_out_dir = train_probe.remote(
@@ -178,6 +187,8 @@ def main(
         out_dir=out_dir,
         epochs=epochs,
         lr=lr,
+        max_examples=max_examples,
+        sample_seed=sample_seed,
         no_chat_template=no_chat_template,
     )
     print(f"Saved probe outputs to Modal Volume path: {remote_out_dir}")

pyproject.toml

@@ -25,6 +25,7 @@ dependencies = [
 train-probe = "scripts.train_probe:main"
 sweep-layers = "scripts.sweep_layers:main"
 guarded-generate = "scripts.guarded_generate:main"
+prepare-training-data = "scripts.prepare_training_data:main"
 
 [tool.setuptools.packages.find]
 include = ["activation_probe_mvp*", "scripts*"]