| """Input classifier: next-token "verdict" training, per paper App. E.7 + E.1. |
| |
| The paper frames harmfulness classification as next-token prediction |
| (predicting "yes"/"no" as the next token after a task template) and shows this |
| beats a value head (App. E.1). Init from the BASE model, not a chat-tuned one |
| (App. E.4: harmlessness biases conflict with the constitution's fine-grained |
| harm definition). |
| |
| Grounded on huggingface/trl `SFTConfig` (main, read 2026-11): prompt-completion |
| datasets + `completion_only_loss=True` give exactly the single-token-target NTP |
| behavior we need. |
| |
| Run: |
| python scripts/train_input_classifier.py \ |
| --train data/train.jsonl --val data/val.jsonl \ |
| --hub_model_id aytsrndl/cc-repro-input-classifier |
| """ |
| import argparse, json, random |
| from datasets import Dataset |
| from trl import SFTTrainer, SFTConfig |
|
|
| random.seed(7) |
|
|
| TASK = ( |
| "Constitution: classify the user's request against the CBRN/chemistry " |
| "harmful constitution (chemical/biological/radiological/nuclear weapons, " |
| "explosives, controlled-substance synthesis, poisoning) AND its harmless " |
| "constitution (legitimate chemistry, dual-use-but-legal topics, everyday " |
| "requests). Harmful = requests actionable harm per the harmful categories. " |
| "Answer with a single word: yes (harmful) or no (harmless).\n\n" |
| "Request: {request}\n\nVerdict:" |
| ) |
|
|
|
|
| def load(path, limit=None): |
| rows = [json.loads(l) for l in open(path)] |
| if limit: |
| rows = rows[:limit] |
| return rows |
|
|
|
|
| def to_sft(rows, keep_system_variants=True): |
| out = [] |
| for r in rows: |
| |
| out.append({"prompt": TASK.format(request=r["user"]), |
| "completion": " yes" if r["label"] == 1 else " no"}) |
| return out |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--base_model", default="Qwen/Qwen2.5-1.5B", |
| help="BASE (pretraining-only) checkpoint, per App. E.4") |
| ap.add_argument("--train", default="data/train.jsonl") |
| ap.add_argument("--val", default="data/val.jsonl") |
| ap.add_argument("--max_length", type=int, default=1024) |
| ap.add_argument("--lr", type=float, default=1e-5) |
| ap.add_argument("--epochs", type=float, default=1.0) |
| ap.add_argument("--per_device_bs", type=int, default=16) |
| ap.add_argument("--grad_accum", type=int, default=2) |
| ap.add_argument("--hub_model_id", required=True) |
| args = ap.parse_args() |
|
|
| train_ds = Dataset.from_list(to_sft(load(args.train))) |
| val_ds = Dataset.from_list(to_sft(load(args.val, limit=2000))) |
| assert "prompt" in train_ds.column_names and "completion" in train_ds.column_names |
|
|
| cfg = SFTConfig( |
| output_dir="out/input-classifier", |
| hub_model_id=args.hub_model_id, |
| push_to_hub=True, |
| completion_only_loss=True, |
| max_length=args.max_length, |
| packing=False, |
| num_train_epochs=args.epochs, |
| learning_rate=args.lr, |
| per_device_train_batch_size=args.per_device_bs, |
| per_device_eval_batch_size=args.per_device_bs, |
| gradient_accumulation_steps=args.grad_accum, |
| eval_strategy="steps", |
| eval_steps=100, |
| logging_steps=10, |
| save_strategy="steps", |
| save_steps=500, |
| report_to=["trackio"], |
| seed=7, |
| ) |
| trainer = SFTTrainer(model=args.base_model, args=cfg, |
| train_dataset=train_ds, eval_dataset=val_ds) |
| trainer.train() |
| metrics = trainer.evaluate() |
| print("eval:", metrics) |
| trainer.push_to_hub(commit_message="input classifier (NTP verdict, base-model init)") |
| print(f"pushed -> https://huggingface.co/{args.hub_model_id}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
