constitutional-classifiers-repro / scripts /train_output_classifier.py

output classifier trainer (value head, omega ramp)

808f43d verified 1 day ago

6.98 kB

"""Output classifier: streaming harmfulness head, per paper §3.4.
A value head scores harmfulness of the FULL sequence from every prefix, so a
streaming guard can cut generation mid-stream. Loss (paper Eq. 1):
    L = lam * L_NTP(t_1:T)  +  sum_{T'=1..T} BCE(y, p(y=1 | t_1:T'))
    p = (1-omega) * sigmoid(z_{T'})  +  omega * max_{tau<=T'} sigmoid(z_tau)
omega ramps linearly 0 -> 1, reaching 1 at 3/4 of training (optimization
stability). lam and exact LR are NOT published in the paper; defaults here are
our choice and are flagged in the README.
Run:
  python scripts/train_output_classifier.py \
      --base_model Qwen/Qwen2.5-1.5B \
      --train data/train.jsonl --val data/val.jsonl \
      --hub_model_id aytsrndl/cc-repro-output-classifier
"""
import argparse, json, math, os, torch, torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset
from transformers import (AutoModelForCausalLM, AutoTokenizer, Trainer,
                          TrainingArguments, TrainerCallback)


class StreamClassifier(nn.Module):
    """Causal LM + linear score head; forward returns the paper's loss."""

    def __init__(self, model_name):
        super().__init__()
        self.lm = AutoModelForCausalLM.from_pretrained(model_name)
        self.score_head = nn.Linear(self.lm.config.hidden_size, 1)
        self.omega = 0.0   # updated each step by OmegaRamp
        self.lam = None    # set from args in the script

    def forward(self, input_ids, attention_mask, ntp_labels, harm_labels):
        out = self.lm(input_ids=input_ids, attention_mask=attention_mask,
                      labels=ntp_labels, output_hidden_states=True)
        h = out.hidden_states[-1]                        # (B, T, H) last-layer states
        z = self.score_head(h.float()).squeeze(-1)       # (B, T) harmfulness logits
        p_tok = torch.sigmoid(z)
        cummax, _ = torch.cummax(p_tok, dim=1)           # max_{tau<=T'} sigma(z_tau)
        p = (1 - self.omega) * p_tok + self.omega * cummax
        y = harm_labels.unsqueeze(1).expand_as(p).float()
        bce = F.binary_cross_entropy(p, y, reduction="none")
        harm_loss = (bce * attention_mask).sum() / attention_mask.sum().clamp(min=1)
        ntp_loss = out.loss if out.loss is not None else torch.tensor(0.0, device=z.device)
        loss = self.lam * ntp_loss + harm_loss
        return {"loss": loss, "harm_loss": harm_loss.detach(), "ntp_loss": ntp_loss.detach()}


class PairDataset(Dataset):
    """Conversation = template(user)+assistant; label = 1 if harmful."""

    def __init__(self, path, tokenizer, max_length):
        self.rows = [json.loads(l) for l in open(path)]
        self.tok = tokenizer
        self.max_length = max_length

    def __len__(self):
        return len(self.rows)

    def __getitem__(self, i):
        r = self.rows[i]
        text = f"Request: {r['user']}\n\nAssistant response: {r['assistant']}"
        ids = self.tok(text, truncation=True, max_length=self.max_length)
        return {"input_ids": ids["input_ids"], "harm": float(r["label"])}


def collate(batch, pad_id):
    maxlen = max(len(b["input_ids"]) for b in batch)
    input_ids, attn, ntp, harm = [], [], [], []
    for b in batch:
        n, pad = len(b["input_ids"]), maxlen - len(b["input_ids"])
        input_ids.append(b["input_ids"] + [pad_id] * pad)
        attn.append([1] * n + [0] * pad)
        ntp.append(b["input_ids"] + [-100] * pad)   # NTP regularizer on real tokens
        harm.append(b["harm"])
    return {"input_ids": torch.tensor(input_ids), "attention_mask": torch.tensor(attn),
            "ntp_labels": torch.tensor(ntp), "harm_labels": torch.tensor(harm)}


class OmegaRamp(TrainerCallback):
    """omega: 0 -> 1 linearly, reaching 1 at 3/4 of training."""

    def __init__(self, model, total_steps):
        self.model, self.total = model, total_steps

    def on_step_end(self, args, state, control, **kw):
        frac = state.global_step / max(1, math.floor(0.75 * self.total))
        self.model.omega = min(1.0, frac)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--base_model", default="Qwen/Qwen2.5-1.5B")
    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, help="NOT published in the paper; our default")
    ap.add_argument("--lam", type=float, default=0.5, help="NTP regularizer weight; NOT published")
    ap.add_argument("--epochs", type=float, default=1.0)
    ap.add_argument("--per_device_bs", type=int, default=8)
    ap.add_argument("--grad_accum", type=int, default=4)
    ap.add_argument("--hub_model_id", required=True)
    args = ap.parse_args()

    tok = AutoTokenizer.from_pretrained(args.base_model)
    if tok.pad_token is None:
        tok.pad_token = tok.eos_token
    model = StreamClassifier(args.base_model)
    model.lam = args.lam

    train_ds = PairDataset(args.train, tok, args.max_length)
    val_ds = PairDataset(args.val, tok, args.max_length)
    steps = math.ceil(len(train_ds) / (args.per_device_bs * args.grad_accum) * args.epochs)

    targs = TrainingArguments(
        output_dir="out/output-classifier",
        hub_model_id=args.hub_model_id,
        push_to_hub=False,  # custom class: saved manually below
        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,
        gradient_checkpointing=True,
        bf16=True,
        eval_strategy="steps",
        eval_steps=100,
        logging_steps=10,
        save_strategy="no",
        remove_unused_columns=False,   # our collate emits nonstandard keys
        report_to=["trackio"],
        seed=7,
    )
    trainer = Trainer(model=model, args=targs, train_dataset=train_ds, eval_dataset=val_ds,
                      data_collator=lambda b: collate(b, tok.pad_token_id),
                      callbacks=[OmegaRamp(model, steps)])
    trainer.train()
    metrics = trainer.evaluate()
    print("eval:", metrics)

    # Persist: LM weights + score head + loader, then push (custom class can't be
    # auto-loaded by AutoModel, so we ship the pieces).
    os.makedirs("out/final", exist_ok=True)
    model.lm.save_pretrained("out/final")
    tok.save_pretrained("out/final")
    torch.save(model.score_head.state_dict(), "out/final/score_head.pt")
    with open("out/final/loss_config.json", "w") as f:
        json.dump({"lam": args.lam, "omega_end": 1.0, "paper": "arXiv:2501.18837 §3.4"}, f)
    from huggingface_hub import HfApi
    api = HfApi()
    api.create_repo(args.hub_model_id, exist_ok=True)
    api.upload_folder(folder_path="out/final", repo_id=args.hub_model_id)
    print(f"pushed -> https://huggingface.co/{args.hub_model_id}")


if __name__ == "__main__":
    main()

	"""Output classifier: streaming harmfulness head, per paper §3.4.

	A value head scores harmfulness of the FULL sequence from every prefix, so a
	streaming guard can cut generation mid-stream. Loss (paper Eq. 1):
	L = lam * L_NTP(t_1:T) + sum_{T'=1..T} BCE(y, p(y=1 \| t_1:T'))
	p = (1-omega) * sigmoid(z_{T'}) + omega * max_{tau<=T'} sigmoid(z_tau)
	omega ramps linearly 0 -> 1, reaching 1 at 3/4 of training (optimization
	stability). lam and exact LR are NOT published in the paper; defaults here are
	our choice and are flagged in the README.

	Run:
	python scripts/train_output_classifier.py \
	--base_model Qwen/Qwen2.5-1.5B \
	--train data/train.jsonl --val data/val.jsonl \
	--hub_model_id aytsrndl/cc-repro-output-classifier
	"""
	import argparse, json, math, os, torch, torch.nn as nn
	import torch.nn.functional as F
	from torch.utils.data import Dataset
	from transformers import (AutoModelForCausalLM, AutoTokenizer, Trainer,
	TrainingArguments, TrainerCallback)


	class StreamClassifier(nn.Module):
	"""Causal LM + linear score head; forward returns the paper's loss."""

	def __init__(self, model_name):
	super().__init__()
	self.lm = AutoModelForCausalLM.from_pretrained(model_name)
	self.score_head = nn.Linear(self.lm.config.hidden_size, 1)
	self.omega = 0.0 # updated each step by OmegaRamp
	self.lam = None # set from args in the script

	def forward(self, input_ids, attention_mask, ntp_labels, harm_labels):
	out = self.lm(input_ids=input_ids, attention_mask=attention_mask,
	labels=ntp_labels, output_hidden_states=True)
	h = out.hidden_states[-1] # (B, T, H) last-layer states
	z = self.score_head(h.float()).squeeze(-1) # (B, T) harmfulness logits
	p_tok = torch.sigmoid(z)
	cummax, _ = torch.cummax(p_tok, dim=1) # max_{tau<=T'} sigma(z_tau)
	p = (1 - self.omega) * p_tok + self.omega * cummax
	y = harm_labels.unsqueeze(1).expand_as(p).float()
	bce = F.binary_cross_entropy(p, y, reduction="none")
	harm_loss = (bce * attention_mask).sum() / attention_mask.sum().clamp(min=1)
	ntp_loss = out.loss if out.loss is not None else torch.tensor(0.0, device=z.device)
	loss = self.lam * ntp_loss + harm_loss
	return {"loss": loss, "harm_loss": harm_loss.detach(), "ntp_loss": ntp_loss.detach()}


	class PairDataset(Dataset):
	"""Conversation = template(user)+assistant; label = 1 if harmful."""

	def __init__(self, path, tokenizer, max_length):
	self.rows = [json.loads(l) for l in open(path)]
	self.tok = tokenizer
	self.max_length = max_length

	def __len__(self):
	return len(self.rows)

	def __getitem__(self, i):
	r = self.rows[i]
	text = f"Request: {r['user']}\n\nAssistant response: {r['assistant']}"
	ids = self.tok(text, truncation=True, max_length=self.max_length)
	return {"input_ids": ids["input_ids"], "harm": float(r["label"])}


	def collate(batch, pad_id):
	maxlen = max(len(b["input_ids"]) for b in batch)
	input_ids, attn, ntp, harm = [], [], [], []
	for b in batch:
	n, pad = len(b["input_ids"]), maxlen - len(b["input_ids"])
	input_ids.append(b["input_ids"] + [pad_id] * pad)
	attn.append([1] * n + [0] * pad)
	ntp.append(b["input_ids"] + [-100] * pad) # NTP regularizer on real tokens
	harm.append(b["harm"])
	return {"input_ids": torch.tensor(input_ids), "attention_mask": torch.tensor(attn),
	"ntp_labels": torch.tensor(ntp), "harm_labels": torch.tensor(harm)}


	class OmegaRamp(TrainerCallback):
	"""omega: 0 -> 1 linearly, reaching 1 at 3/4 of training."""

	def __init__(self, model, total_steps):
	self.model, self.total = model, total_steps

	def on_step_end(self, args, state, control, **kw):
	frac = state.global_step / max(1, math.floor(0.75 * self.total))
	self.model.omega = min(1.0, frac)


	def main():
	ap = argparse.ArgumentParser()
	ap.add_argument("--base_model", default="Qwen/Qwen2.5-1.5B")
	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, help="NOT published in the paper; our default")
	ap.add_argument("--lam", type=float, default=0.5, help="NTP regularizer weight; NOT published")
	ap.add_argument("--epochs", type=float, default=1.0)
	ap.add_argument("--per_device_bs", type=int, default=8)
	ap.add_argument("--grad_accum", type=int, default=4)
	ap.add_argument("--hub_model_id", required=True)
	args = ap.parse_args()

	tok = AutoTokenizer.from_pretrained(args.base_model)
	if tok.pad_token is None:
	tok.pad_token = tok.eos_token
	model = StreamClassifier(args.base_model)
	model.lam = args.lam

	train_ds = PairDataset(args.train, tok, args.max_length)
	val_ds = PairDataset(args.val, tok, args.max_length)
	steps = math.ceil(len(train_ds) / (args.per_device_bs * args.grad_accum) * args.epochs)

	targs = TrainingArguments(
	output_dir="out/output-classifier",
	hub_model_id=args.hub_model_id,
	push_to_hub=False, # custom class: saved manually below
	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,
	gradient_checkpointing=True,
	bf16=True,
	eval_strategy="steps",
	eval_steps=100,
	logging_steps=10,
	save_strategy="no",
	remove_unused_columns=False, # our collate emits nonstandard keys
	report_to=["trackio"],
	seed=7,
	)
	trainer = Trainer(model=model, args=targs, train_dataset=train_ds, eval_dataset=val_ds,
	data_collator=lambda b: collate(b, tok.pad_token_id),
	callbacks=[OmegaRamp(model, steps)])
	trainer.train()
	metrics = trainer.evaluate()
	print("eval:", metrics)

	# Persist: LM weights + score head + loader, then push (custom class can't be
	# auto-loaded by AutoModel, so we ship the pieces).
	os.makedirs("out/final", exist_ok=True)
	model.lm.save_pretrained("out/final")
	tok.save_pretrained("out/final")
	torch.save(model.score_head.state_dict(), "out/final/score_head.pt")
	with open("out/final/loss_config.json", "w") as f:
	json.dump({"lam": args.lam, "omega_end": 1.0, "paper": "arXiv:2501.18837 §3.4"}, f)
	from huggingface_hub import HfApi
	api = HfApi()
	api.create_repo(args.hub_model_id, exist_ok=True)
	api.upload_folder(folder_path="out/final", repo_id=args.hub_model_id)
	print(f"pushed -> https://huggingface.co/{args.hub_model_id}")


	if __name__ == "__main__":
	main()