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