diff --git a/.gitattributes b/.gitattributes new file mode 100644 index 0000000..9fdd07c --- /dev/null +++ b/.gitattributes @@ -0,0 +1 @@ +data/training_data.jsonl filter=lfs diff=lfs merge=lfs -text diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 6bf1ce0..1e32464 100644 --- a/.github/workflows/ci.yml +++ b/.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 diff --git a/README.md b/README.md index 62deb8d..db4d5ad 100644 --- a/README.md +++ b/README.md @@ -31,9 +31,9 @@ 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: @@ -41,8 +41,8 @@ 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,9 +229,9 @@ 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: @@ -203,8 +239,8 @@ 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" ``` diff --git a/activation_probe_mvp/data.py b/activation_probe_mvp/data.py index 9401831..cf43a47 100644 --- a/activation_probe_mvp/data.py +++ b/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. diff --git a/data/README.md b/data/README.md new file mode 100644 index 0000000..22dc07f --- /dev/null +++ b/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` diff --git a/data/dataset_summary.json b/data/dataset_summary.json new file mode 100644 index 0000000..d208565 --- /dev/null +++ b/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" +} diff --git a/data/training_data.jsonl b/data/training_data.jsonl new file mode 100644 index 0000000..7637a11 --- /dev/null +++ b/data/training_data.jsonl @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:34c1a80368687863cc18a3631e8e1b19f4e651da995598098fbec06011450c30 +size 211058673 diff --git a/modal_train.py b/modal_train.py index 8829724..7c1af9a 100644 --- a/modal_train.py +++ b/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}") diff --git a/pyproject.toml b/pyproject.toml index c4eebe8..8e377f7 100644 --- a/pyproject.toml +++ b/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*"] diff --git a/scripts/prepare_training_data.py b/scripts/prepare_training_data.py new file mode 100644 index 0000000..92b70cc --- /dev/null +++ b/scripts/prepare_training_data.py @@ -0,0 +1,398 @@ +#!/usr/bin/env python +from __future__ import annotations + +import argparse +import hashlib +import json +import re +from collections import Counter, defaultdict +from collections.abc import Iterable +from pathlib import Path +from typing import Any + +from datasets import DatasetDict, load_dataset + + +DEFAULT_PAIRWISE_DATASET_ID = "PKU-Alignment/PKU-SafeRLHF" +DEFAULT_SINGLE_RESPONSE_DATASET_ID = "PKU-Alignment/BeaverTails" +DEFAULT_DATASET_VERSION = "public_safety_v1" +DEFAULT_SINGLE_RESPONSE_MAX_RECORDS = 100_000 +SCHEMA_FIELDS = [ + "id", + "prompt", + "response", + "label", + "split", + "topic", + "prompt_intent", + "response_behavior", + "dataset_version", +] + + +def normalize_split(split: str, *, prompt: str, validation_fraction: float) -> str: + split = split.lower() + if split in {"test", "eval", "evaluation"} or split.endswith("_test"): + return "test" + if split in {"dev", "val", "valid", "validation"} or split.endswith("_validation"): + return "validation" + if validation_fraction <= 0: + return "train" + + # Split by prompt so response_0 and response_1 never leak across splits. + digest = hashlib.sha256(prompt.encode("utf-8")).hexdigest() + bucket = int(digest[:8], 16) / 0xFFFFFFFF + return "validation" if bucket < validation_fraction else "train" + + +def normalize_topic(value: str) -> str: + text = value.strip().lower() + text = re.sub(r"[^a-z0-9]+", "_", text).strip("_") + return text or "unknown" + + +def parse_harm_categories(value: Any) -> dict[str, bool]: + if value is None: + return {} + if isinstance(value, str): + value = json.loads(value) + if not isinstance(value, dict): + return {} + return {str(key): bool(val) for key, val in value.items()} + + +def topic_from_categories(value: Any) -> str: + categories = parse_harm_categories(value) + active = [name for name, enabled in categories.items() if enabled] + if not active: + return "safe" + return normalize_topic(active[0]) + + +def stable_key(record: dict[str, Any]) -> str: + return hashlib.sha256(f"{record['prompt']}\0{record['response']}".encode("utf-8")).hexdigest() + + +def source_from_id(record_id: str) -> str: + parts = record_id.split(":") + return parts[1] if len(parts) >= 4 else "unknown" + + +def convert_pairwise_response( + row: dict[str, Any], + *, + response_idx: int, + source_split: str, + record_index: int, + dataset_version: str, + source_name: str, + validation_fraction: float, +) -> dict[str, Any] | None: + prompt = row.get("prompt") + response = row.get(f"response_{response_idx}") + is_safe = row.get(f"is_response_{response_idx}_safe") + if not prompt or not response or not isinstance(is_safe, bool): + return None + + label = 0 if is_safe else 1 + split = normalize_split( + source_split, + prompt=str(prompt), + validation_fraction=validation_fraction, + ) + topic = topic_from_categories(row.get(f"response_{response_idx}_harm_category")) + if label == 1 and topic == "safe": + topic = "policy_violation" + + return { + "id": f"{dataset_version}:{source_name}:{split}:{record_index:06d}", + "prompt": str(prompt), + "response": str(response), + "label": label, + "split": split, + "topic": topic, + "prompt_intent": "unknown", + "response_behavior": "allowed_response" if label == 0 else "policy_violation", + "dataset_version": dataset_version, + } + + +def convert_single_response_row( + row: dict[str, Any], + *, + source_split: str, + record_index: int, + dataset_version: str, + validation_fraction: float, +) -> dict[str, Any] | None: + prompt = row.get("prompt") + response = row.get("response") + is_safe = row.get("is_safe") + if not prompt or not response or not isinstance(is_safe, bool): + return None + + label = 0 if is_safe else 1 + split = normalize_split( + source_split, + prompt=str(prompt), + validation_fraction=validation_fraction, + ) + topic = topic_from_categories(row.get("category")) + if label == 1 and topic == "safe": + topic = "policy_violation" + + return { + "id": f"{dataset_version}:single_response_safety:{split}:{record_index:06d}", + "prompt": str(prompt), + "response": str(response), + "label": label, + "split": split, + "topic": topic, + "prompt_intent": "unknown", + "response_behavior": "allowed_response" if label == 0 else "policy_violation", + "dataset_version": dataset_version, + } + + +def iter_pairwise_records( + dataset: DatasetDict, + *, + dataset_version: str, + dedupe: bool, + splits: Iterable[str] | None, + validation_fraction: float, +) -> list[dict[str, Any]]: + records: list[dict[str, Any]] = [] + seen_pairs: set[tuple[str, str]] = set() + requested_splits = None if splits is None else set(splits) + + for source_split, split_dataset in dataset.items(): + if requested_splits is not None and source_split not in requested_splits: + continue + for row in split_dataset: + for response_idx in (0, 1): + record = convert_pairwise_response( + dict(row), + response_idx=response_idx, + source_split=source_split, + record_index=len(records), + dataset_version=dataset_version, + source_name="pairwise_safety", + validation_fraction=validation_fraction, + ) + if record is None: + continue + pair = (record["prompt"], record["response"]) + if dedupe and pair in seen_pairs: + continue + seen_pairs.add(pair) + records.append(record) + + return records + + +def iter_single_response_records( + dataset: DatasetDict, + *, + dataset_version: str, + dedupe: bool, + splits: Iterable[str] | None, + validation_fraction: float, +) -> list[dict[str, Any]]: + records: list[dict[str, Any]] = [] + seen_pairs: set[tuple[str, str]] = set() + requested_splits = None if splits is None else set(splits) + + for source_split, split_dataset in dataset.items(): + if requested_splits is not None and source_split not in requested_splits: + continue + for row in split_dataset: + record = convert_single_response_row( + dict(row), + source_split=source_split, + record_index=len(records), + dataset_version=dataset_version, + validation_fraction=validation_fraction, + ) + if record is None: + continue + pair = (record["prompt"], record["response"]) + if dedupe and pair in seen_pairs: + continue + seen_pairs.add(pair) + records.append(record) + + return records + + +def cap_records_diverse( + records: list[dict[str, Any]], max_records: int | None +) -> list[dict[str, Any]]: + if max_records is None or max_records <= 0 or len(records) <= max_records: + return records + + grouped: dict[tuple[str, int, str], list[dict[str, Any]]] = defaultdict(list) + for record in records: + grouped[(record["split"], record["label"], record["topic"])].append(record) + + for group_records in grouped.values(): + group_records.sort(key=stable_key) + + selected: list[dict[str, Any]] = [] + group_keys = sorted(grouped) + while len(selected) < max_records and group_keys: + next_keys: list[tuple[str, int, str]] = [] + for key in group_keys: + if grouped[key]: + selected.append(grouped[key].pop(0)) + if len(selected) == max_records: + break + if grouped[key]: + next_keys.append(key) + group_keys = next_keys + + return selected + + +def dedupe_records(records: list[dict[str, Any]]) -> list[dict[str, Any]]: + deduped: list[dict[str, Any]] = [] + seen_pairs: set[tuple[str, str]] = set() + for record in records: + pair = (record["prompt"], record["response"]) + if pair in seen_pairs: + continue + seen_pairs.add(pair) + deduped.append(record) + return deduped + + +def write_jsonl(path: Path, records: list[dict[str, Any]]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("w", encoding="utf-8") as f: + for record in records: + f.write(json.dumps(record, ensure_ascii=False, separators=(",", ":")) + "\n") + + +def build_summary(records: list[dict[str, Any]], output_path: Path) -> dict[str, Any]: + return { + "file": str(output_path), + "total_rows": len(records), + "source_distribution": dict( + sorted(Counter(source_from_id(r["id"]) for r in records).items()) + ), + "label_distribution": dict(sorted(Counter(str(r["label"]) for r in records).items())), + "split_distribution": dict(sorted(Counter(r["split"] for r in records).items())), + "behavior_distribution": dict( + sorted(Counter(r["response_behavior"] for r in records).items()) + ), + "prompt_intent_distribution": dict( + sorted(Counter(r["prompt_intent"] for r in records).items()) + ), + "topic_distribution": dict(sorted(Counter(r["topic"] for r in records).items())), + "topic_count": len({r["topic"] for r in records}), + "unique_ids": len({r["id"] for r in records}), + "unique_prompts": len({r["prompt"] for r in records}), + "unique_prompt_response_pairs": len({(r["prompt"], r["response"]) for r in records}), + "schema_fields": SCHEMA_FIELDS, + "dataset_version": records[0]["dataset_version"] if records else None, + } + + +def main() -> None: + parser = argparse.ArgumentParser( + description="Convert public safety datasets into Open Constitution JSONL." + ) + parser.add_argument("--pairwise_dataset_id", default=DEFAULT_PAIRWISE_DATASET_ID) + parser.add_argument("--single_response_dataset_id", default=DEFAULT_SINGLE_RESPONSE_DATASET_ID) + parser.add_argument("--output_path", default="data/training_data.jsonl") + parser.add_argument("--summary_path", default="data/dataset_summary.json") + parser.add_argument("--dataset_version", default=DEFAULT_DATASET_VERSION) + parser.add_argument( + "--single_response_max_records", + type=int, + default=DEFAULT_SINGLE_RESPONSE_MAX_RECORDS, + help=( + "Maximum single-response rows to include after deterministic label/topic-diverse " + "sampling. Use 0 to include all rows." + ), + ) + parser.add_argument( + "--no_single_response", + action="store_true", + help="Only convert the pairwise dataset.", + ) + parser.add_argument( + "--validation_fraction", + type=float, + default=0.1, + help="Fraction of source train prompts assigned to validation.", + ) + parser.add_argument( + "--pairwise_splits", + nargs="*", + default=["train", "test"], + help="Pairwise source splits to convert. Use all to convert every split.", + ) + parser.add_argument( + "--single_response_splits", + nargs="*", + default=["330k_train", "330k_test"], + help="Single-response source splits to convert. Use all to convert every split.", + ) + parser.add_argument( + "--keep_duplicates", + action="store_true", + help="Keep duplicate prompt/response pairs instead of dropping them.", + ) + args = parser.parse_args() + + if not 0 <= args.validation_fraction < 1: + raise ValueError("--validation_fraction must be in [0, 1)") + + output_path = Path(args.output_path) + summary_path = Path(args.summary_path) + + pairwise_dataset = load_dataset(args.pairwise_dataset_id) + if not isinstance(pairwise_dataset, DatasetDict): + pairwise_dataset = DatasetDict({"train": pairwise_dataset}) + records = iter_pairwise_records( + pairwise_dataset, + dataset_version=args.dataset_version, + dedupe=not args.keep_duplicates, + splits=None if args.pairwise_splits == ["all"] else args.pairwise_splits, + validation_fraction=args.validation_fraction, + ) + + if not args.no_single_response: + single_response_dataset = load_dataset(args.single_response_dataset_id) + if not isinstance(single_response_dataset, DatasetDict): + single_response_dataset = DatasetDict({"train": single_response_dataset}) + single_response_records = iter_single_response_records( + single_response_dataset, + dataset_version=args.dataset_version, + dedupe=not args.keep_duplicates, + splits=None if args.single_response_splits == ["all"] else args.single_response_splits, + validation_fraction=args.validation_fraction, + ) + records.extend( + cap_records_diverse( + single_response_records, + None if args.single_response_max_records == 0 else args.single_response_max_records, + ) + ) + + if not args.keep_duplicates: + records = dedupe_records(records) + + if not records: + raise ValueError("No usable prompt/response rows were found.") + + write_jsonl(output_path, records) + summary = build_summary(records, output_path) + summary_path.write_text(json.dumps(summary, indent=2) + "\n", encoding="utf-8") + + print(json.dumps(summary, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/scripts/sweep_layers.py b/scripts/sweep_layers.py index 3dd318b..3b20067 100644 --- a/scripts/sweep_layers.py +++ b/scripts/sweep_layers.py @@ -19,7 +19,7 @@ def main(): parser = argparse.ArgumentParser() parser.add_argument("--model_id", default="google/gemma-4-E2B-it") - parser.add_argument("--data_path", default="data/examples.jsonl") + parser.add_argument("--data_path", default="data/training_data.jsonl") parser.add_argument("--layers", default="-2,-4,-6,-8,-10,-12") parser.add_argument("--epochs", type=int, default=60) parser.add_argument( diff --git a/scripts/train_probe.py b/scripts/train_probe.py index 7be3797..59217fd 100644 --- a/scripts/train_probe.py +++ b/scripts/train_probe.py @@ -12,7 +12,7 @@ load_model_and_tokenizer, ) from activation_probe_mvp.chat import format_exchange -from activation_probe_mvp.data import load_jsonl +from activation_probe_mvp.data import balanced_sample, load_jsonl from activation_probe_mvp.training import evaluate_probe, save_probe, train_linear_probe @@ -23,11 +23,18 @@ def main(): default="google/gemma-4-E2B-it", help="Default is Gemma 4 E2B instruction-tuned. Also works with Qwen/Llama-style causal LMs.", ) - parser.add_argument("--data_path", default="data/examples.jsonl") + parser.add_argument("--data_path", default="data/training_data.jsonl") parser.add_argument("--layer", type=int, default=-4) - parser.add_argument("--out_dir", default="./probe_out") + parser.add_argument("--out_dir", default="./probe_out_public_safety") parser.add_argument("--epochs", type=int, default=100) parser.add_argument("--lr", type=float, default=1e-3) + parser.add_argument( + "--max_examples", + type=int, + default=20000, + help="Maximum examples to use after deterministic label-balanced sampling. Use 0 for all.", + ) + parser.add_argument("--sample_seed", type=int, default=0) parser.add_argument( "--no_chat_template", action="store_true", @@ -38,7 +45,13 @@ def main(): print(f"Loading model: {args.model_id}") model, tokenizer_or_processor, device = load_model_and_tokenizer(args.model_id) - examples = load_jsonl(args.data_path) + all_examples = load_jsonl(args.data_path) + examples = balanced_sample( + all_examples, + max_examples=args.max_examples if args.max_examples > 0 else None, + seed=args.sample_seed, + ) + print(f"Loaded {len(all_examples)} examples; training on {len(examples)}") if len(examples) < 4: raise ValueError("Need at least a few examples. For real use, use thousands.") diff --git a/tests/test_data.py b/tests/test_data.py index 0e1d724..8664067 100644 --- a/tests/test_data.py +++ b/tests/test_data.py @@ -2,7 +2,7 @@ import pytest -from activation_probe_mvp.data import Example, load_jsonl +from activation_probe_mvp.data import Example, balanced_sample, load_jsonl def test_load_jsonl_reads_examples_and_skips_blank_lines(tmp_path): @@ -36,3 +36,25 @@ def test_load_jsonl_rejects_unknown_labels(tmp_path): with pytest.raises(ValueError, match="label must be 0 or 1"): load_jsonl(data_path) + + +def test_balanced_sample_is_deterministic_and_balances_labels(): + examples = [Example(prompt=f"safe {i}", response="ok", label=0) for i in range(10)] + [ + Example(prompt=f"unsafe {i}", response="bad", label=1) for i in range(10) + ] + + sample_a = balanced_sample(examples, max_examples=6, seed=123) + sample_b = balanced_sample(examples, max_examples=6, seed=123) + + assert sample_a == sample_b + assert len(sample_a) == 6 + assert sum(example.label == 0 for example in sample_a) == 3 + assert sum(example.label == 1 for example in sample_a) == 3 + + +def test_balanced_sample_returns_all_examples_when_not_capped(): + examples = [Example(prompt="hello", response="hi", label=0)] + + assert balanced_sample(examples, max_examples=None) == examples + assert balanced_sample(examples, max_examples=0) == examples + assert balanced_sample(examples, max_examples=10) == examples diff --git a/tests/test_handcrafted_dataset.py b/tests/test_handcrafted_dataset.py new file mode 100644 index 0000000..59b6a52 --- /dev/null +++ b/tests/test_handcrafted_dataset.py @@ -0,0 +1,95 @@ +from __future__ import annotations + +import json +from collections import Counter +from pathlib import Path + +from activation_probe_mvp.data import load_jsonl + + +DATASET_PATH = Path("data/training_data.jsonl") +SUMMARY_PATH = Path("data/dataset_summary.json") +REQUIRED_FIELDS = { + "id", + "prompt", + "response", + "label", + "split", + "topic", + "prompt_intent", + "response_behavior", + "dataset_version", +} + + +def load_records() -> list[dict]: + with DATASET_PATH.open("r", encoding="utf-8") as f: + return [json.loads(line) for line in f if line.strip()] + + +def load_summary() -> dict: + return json.loads(SUMMARY_PATH.read_text(encoding="utf-8")) + + +def source_from_id(record_id: str) -> str: + parts = record_id.split(":") + return parts[1] if len(parts) >= 4 else "unknown" + + +def test_training_dataset_shape_and_metadata(): + records = load_records() + + assert len(records) > 0 + assert all(REQUIRED_FIELDS <= record.keys() for record in records) + assert {record["label"] for record in records} == {0, 1} + assert len({record["id"] for record in records}) == len(records) + + +def test_training_dataset_distribution_metadata_is_present(): + records = load_records() + + assert Counter(record["split"] for record in records) + assert Counter(record["response_behavior"] for record in records) + assert Counter(record["topic"] for record in records) + assert Counter(record["dataset_version"] for record in records) + + +def test_handcrafted_dataset_has_unique_prompt_response_pairs(): + records = load_records() + pairs = {(record["prompt"], record["response"]) for record in records} + + assert len(pairs) == len(records) + + +def test_handcrafted_dataset_is_compatible_with_current_loader(): + examples = load_jsonl(DATASET_PATH) + + assert len(examples) == len(load_records()) + assert {example.label for example in examples} == {0, 1} + + +def test_dataset_summary_matches_training_data(): + records = load_records() + summary = load_summary() + + assert summary["file"] == str(DATASET_PATH) + assert summary["total_rows"] == len(records) + assert summary["source_distribution"] == dict( + sorted(Counter(source_from_id(record["id"]) for record in records).items()) + ) + assert summary["label_distribution"] == dict( + sorted(Counter(str(record["label"]) for record in records).items()) + ) + assert summary["split_distribution"] == dict( + sorted(Counter(record["split"] for record in records).items()) + ) + assert summary["behavior_distribution"] == dict( + sorted(Counter(record["response_behavior"] for record in records).items()) + ) + assert summary["topic_distribution"] == dict( + sorted(Counter(record["topic"] for record in records).items()) + ) + assert summary["unique_ids"] == len({record["id"] for record in records}) + assert summary["unique_prompt_response_pairs"] == len( + {(record["prompt"], record["response"]) for record in records} + ) diff --git a/tests/test_prepare_training_data.py b/tests/test_prepare_training_data.py new file mode 100644 index 0000000..8608cdb --- /dev/null +++ b/tests/test_prepare_training_data.py @@ -0,0 +1,260 @@ +from __future__ import annotations + +from scripts.prepare_training_data import ( + build_summary, + cap_records_diverse, + convert_pairwise_response, + convert_single_response_row, + dedupe_records, + iter_pairwise_records, + normalize_split, + topic_from_categories, +) + + +def test_convert_response_maps_safe_response_to_label_0(): + record = convert_pairwise_response( + { + "prompt": "How do I set up a webcam?", + "response_0": "Connect the camera and test your video app.", + "is_response_0_safe": True, + "response_0_harm_category": {"Cybercrime": False}, + }, + response_idx=0, + source_split="train", + record_index=4, + dataset_version="test_public_safety", + source_name="pairwise_safety", + validation_fraction=0, + ) + + assert record == { + "id": "test_public_safety:pairwise_safety:train:000004", + "prompt": "How do I set up a webcam?", + "response": "Connect the camera and test your video app.", + "label": 0, + "split": "train", + "topic": "safe", + "prompt_intent": "unknown", + "response_behavior": "allowed_response", + "dataset_version": "test_public_safety", + } + + +def test_convert_response_maps_unsafe_response_to_label_1_and_topic(): + record = convert_pairwise_response( + { + "prompt": "How do I hide a bribe?", + "response_1": "Use intermediaries and offshore accounts.", + "is_response_1_safe": False, + "response_1_harm_category": { + "Economic Crime": True, + "White-Collar Crime": True, + }, + }, + response_idx=1, + source_split="test", + record_index=8, + dataset_version="test_public_safety", + source_name="pairwise_safety", + validation_fraction=0.1, + ) + + assert record is not None + assert record["label"] == 1 + assert record["split"] == "test" + assert record["topic"] == "economic_crime" + assert record["response_behavior"] == "policy_violation" + + +def test_convert_response_uses_policy_violation_topic_for_uncategorized_unsafe_response(): + record = convert_pairwise_response( + { + "prompt": "Give harmful advice.", + "response_0": "Here is harmful advice.", + "is_response_0_safe": False, + "response_0_harm_category": {"Cybercrime": False}, + }, + response_idx=0, + source_split="train", + record_index=1, + dataset_version="test_public_safety", + source_name="pairwise_safety", + validation_fraction=0, + ) + + assert record is not None + assert record["label"] == 1 + assert record["topic"] == "policy_violation" + + +def test_topic_from_categories_handles_json_string_and_safe_case(): + assert topic_from_categories('{"Cybercrime": true, "Violence": false}') == "cybercrime" + assert topic_from_categories({"Cybercrime": False}) == "safe" + + +def test_normalize_split_keeps_same_prompt_in_same_split(): + prompt = "The same prompt with two candidate responses." + + assert normalize_split("test", prompt=prompt, validation_fraction=0.5) == "test" + assert normalize_split("train", prompt=prompt, validation_fraction=1.0) == "validation" + assert normalize_split("train", prompt=prompt, validation_fraction=0.0) == "train" + + +def test_convert_single_response_row_maps_single_response_schema(): + record = convert_single_response_row( + { + "prompt": "How can I steal from work?", + "response": "Here is a theft workflow.", + "is_safe": False, + "category": { + "financial_crime,property_crime,theft": True, + "privacy_violation": False, + }, + }, + source_split="330k_test", + record_index=5, + dataset_version="test_public_safety", + validation_fraction=0.1, + ) + + assert record is not None + assert record["id"] == "test_public_safety:single_response_safety:test:000005" + assert record["label"] == 1 + assert record["split"] == "test" + assert record["topic"] == "financial_crime_property_crime_theft" + + +def test_iter_pairwise_records_emits_two_responses_and_dedupes_pairs(): + dataset = { + "train": [ + { + "prompt": "same", + "response_0": "safe response", + "response_1": "safe response", + "is_response_0_safe": True, + "is_response_1_safe": True, + "response_0_harm_category": {"Cybercrime": False}, + "response_1_harm_category": {"Cybercrime": False}, + }, + { + "prompt": "unsafe", + "response_0": "bad response", + "response_1": "better response", + "is_response_0_safe": False, + "is_response_1_safe": True, + "response_0_harm_category": {"Cybercrime": True}, + "response_1_harm_category": {"Cybercrime": False}, + }, + ], + "ignored": [ + { + "prompt": "skip", + "response_0": "skip", + "response_1": "skip", + "is_response_0_safe": True, + "is_response_1_safe": True, + } + ], + } + + records = iter_pairwise_records( + dataset, + dataset_version="test_public_safety", + dedupe=True, + splits=["train"], + validation_fraction=0, + ) + + assert [(record["prompt"], record["response"], record["label"]) for record in records] == [ + ("same", "safe response", 0), + ("unsafe", "bad response", 1), + ("unsafe", "better response", 0), + ] + + +def test_cap_records_diverse_samples_across_split_label_topic_groups(): + records = [] + for i in range(10): + records.append( + { + "id": f"test_public_safety:single_response_safety:train:{i:06d}", + "prompt": f"prompt {i}", + "response": f"response {i}", + "label": i % 2, + "split": "train", + "topic": "safe" if i % 2 == 0 else "cybercrime", + "prompt_intent": "unknown", + "response_behavior": "allowed_response" if i % 2 == 0 else "policy_violation", + "dataset_version": "test_public_safety", + } + ) + + capped = cap_records_diverse(records, max_records=4) + + assert len(capped) == 4 + assert {record["label"] for record in capped} == {0, 1} + assert {record["topic"] for record in capped} == {"safe", "cybercrime"} + + +def test_dedupe_records_removes_duplicate_prompt_response_pairs(): + records = [ + { + "id": "test_public_safety:pairwise_safety:train:000000", + "prompt": "same", + "response": "same response", + "label": 0, + "split": "train", + "topic": "safe", + "prompt_intent": "unknown", + "response_behavior": "allowed_response", + "dataset_version": "test_public_safety", + }, + { + "id": "test_public_safety:single_response_safety:train:000001", + "prompt": "same", + "response": "same response", + "label": 0, + "split": "train", + "topic": "safe", + "prompt_intent": "unknown", + "response_behavior": "allowed_response", + "dataset_version": "test_public_safety", + }, + ] + + assert dedupe_records(records) == [records[0]] + + +def test_build_summary_includes_source_distribution(tmp_path): + records = [ + { + "id": "test_public_safety:pairwise_safety:train:000000", + "prompt": "a", + "response": "b", + "label": 0, + "split": "train", + "topic": "safe", + "prompt_intent": "unknown", + "response_behavior": "allowed_response", + "dataset_version": "test_public_safety", + }, + { + "id": "test_public_safety:single_response_safety:train:000001", + "prompt": "c", + "response": "d", + "label": 1, + "split": "train", + "topic": "cybercrime", + "prompt_intent": "unknown", + "response_behavior": "policy_violation", + "dataset_version": "test_public_safety", + }, + ] + + summary = build_summary(records, tmp_path / "training_data.jsonl") + + assert summary["source_distribution"] == { + "pairwise_safety": 1, + "single_response_safety": 1, + }