File size: 3,816 Bytes
966092b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 | """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:
# keep augmented system-prompt variants as prefix context; NTP target is one word
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, # loss on the verdict token only (NTP framing)
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()
|