"""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()