Buckets:
| """Shared training machinery for both classifier scripts (train_modernbert.py, train_qwen.py). | |
| Each model script only decides how to load its backbone, which modules get LoRA, which module is the | |
| classification head, and its data collator. Everything else lives here so both models are trained, | |
| checkpointed and measured identically: | |
| * dataset validation (labels, empty text, case/input leakage between splits); | |
| * LoRA/DoRA attachment with the classification head trained in full; | |
| * tokenization in the shared input format (training.common.tokenize_pairs); | |
| * adapter-only checkpoints at 1/4, 2/4, 3/4 and the end of training, each followed by evaluation; | |
| * scikit-learn metrics, a confusion matrix and per-row validation predictions. | |
| ML imports are kept inside functions so `--preflight` runs without torch installed. | |
| """ | |
| import json | |
| import time | |
| from collections import Counter | |
| from pathlib import Path | |
| import yaml | |
| from training.common import ( | |
| INSTRUCTION, | |
| LABEL2ID, | |
| LABELS, | |
| classification_metrics, | |
| confusion_matrix_counts, | |
| read_rows, | |
| tokenize_pairs, | |
| write_rows, | |
| ) | |
| REQUIRED_KEYS = ( | |
| "model_name", "data_dir", "output_dir", "run_name", "wandb_project", "seed", "max_length", | |
| "epochs", "learning_rate", "lr_scheduler_type", "warmup_ratio", "weight_decay", "max_grad_norm", | |
| "batch_size", "gradient_accumulation_steps", "eval_batch_size", "gradient_checkpointing", | |
| "attention_implementation", "lora_rank", "lora_alpha", "lora_dropout", "use_dora", | |
| "checkpoints_per_run", "logging_steps", "dataloader_num_workers", | |
| ) | |
| def load_config(path, epochs=None): | |
| config = yaml.safe_load(Path(path).read_text(encoding="utf-8")) | |
| missing = [key for key in REQUIRED_KEYS if key not in config] | |
| if missing: | |
| raise ValueError(f"Config {path} is missing: {', '.join(missing)}") | |
| if epochs is not None: | |
| config["epochs"] = epochs | |
| return config | |
| def validate_inputs(config): | |
| """Check the config and the dataset before any model is downloaded.""" | |
| if not 0 < config["max_length"] <= 8192: | |
| raise ValueError("The configured classifier training cap must be between 1 and 8192 tokens") | |
| if config["epochs"] <= 0 or config["batch_size"] <= 0 or config["gradient_accumulation_steps"] <= 0: | |
| raise ValueError("Epochs and batch sizes must be positive") | |
| if config.get("lora_rank", 1) <= 0: | |
| raise ValueError("LoRA rank must be positive") | |
| data_dir = Path(config["data_dir"]) | |
| splits = {name: read_rows(data_dir / f"{name}.jsonl") for name in ("train", "validation")} | |
| if (data_dir / "test_ood.jsonl").exists(): | |
| splits["test_ood"] = read_rows(data_dir / "test_ood.jsonl") | |
| cases, inputs = {}, {} | |
| for split, rows in splits.items(): | |
| if split != "test_ood" and not rows: | |
| raise ValueError(f"{split} split is empty") | |
| for row in rows: | |
| if row["label"] not in LABEL2ID or row["labels"] != LABEL2ID[row["label"]]: | |
| raise ValueError("Dataset label ids disagree with the training label map") | |
| if not row["message"].strip() or not row["policy"].strip(): | |
| raise ValueError("Empty message or policy") | |
| if split != "test_ood" and row["held_out"]: | |
| raise ValueError("Held-out domain present in train/validation") | |
| if row["case_id"] in cases and cases[row["case_id"]] != split: | |
| raise ValueError("A case appears in multiple splits") | |
| key = (row["message"], row["policy"]) | |
| if key in inputs and inputs[key] != split: | |
| raise ValueError("Identical input appears in multiple splits") | |
| cases[row["case_id"]], inputs[key] = split, split | |
| if set(row["label"] for row in splits["train"]) != set(LABELS): | |
| raise ValueError("Training requires all three classes") | |
| return splits | |
| def describe_splits(splits): | |
| return { | |
| name: {"pairs": len(rows), "labels": dict(Counter(row["label"] for row in rows))} | |
| for name, rows in splits.items() | |
| } | |
| def classifier_config(base_config, tokenizer): | |
| """Standard HF three-class single-label setup (cross-entropy loss on the classification head).""" | |
| base_config.num_labels = len(LABELS) | |
| base_config.label2id, base_config.id2label = LABEL2ID, dict(enumerate(LABELS)) | |
| base_config.problem_type = "single_label_classification" | |
| base_config.pad_token_id = tokenizer.pad_token_id | |
| return base_config | |
| def attach_lora(model, config, target_modules, head_modules): | |
| """LoRA (or DoRA when use_dora) on the backbone projections; the new head is trained in full and | |
| saved inside the adapter checkpoint through modules_to_save.""" | |
| from peft import LoraConfig, TaskType, get_peft_model | |
| adapter = LoraConfig( | |
| task_type=TaskType.SEQ_CLS, | |
| r=config["lora_rank"], | |
| lora_alpha=config["lora_alpha"], | |
| lora_dropout=config["lora_dropout"], | |
| use_dora=config["use_dora"], | |
| bias="none", | |
| target_modules=target_modules, | |
| modules_to_save=head_modules, | |
| ) | |
| model = get_peft_model(model, adapter) | |
| model.print_trainable_parameters() | |
| return model | |
| def input_recipe(config, tokenizer, base_config): | |
| """Everything needed to rebuild the exact classifier input at prediction time.""" | |
| return { | |
| "model_name": config["model_name"], | |
| "model_revision": getattr(base_config, "_commit_hash", None), | |
| "model_type": base_config.model_type, | |
| "max_length": config["max_length"], | |
| "instruction": INSTRUCTION, | |
| "label2id": LABEL2ID, | |
| "separator": tokenizer.sep_token or "\n\n", | |
| "pad_token_id": tokenizer.pad_token_id, | |
| "attention_implementation": config["attention_implementation"], | |
| "format": "tokenizer(text=instruction+message+optional_blank_line, text_pair=policy)", | |
| } | |
| def save_adapter(model, tokenizer, recipe, path): | |
| """Adapter weights + trained classification head only; no optimizer or full model weights.""" | |
| path = Path(path) | |
| model.save_pretrained(path) | |
| tokenizer.save_pretrained(path) | |
| (path / "input_format.json").write_text(json.dumps(recipe, indent=2), encoding="utf-8") | |
| def quarter_checkpoint_callback(tokenizer, recipe, output, count): | |
| """Save the adapter and run evaluation at k/count of the total optimizer steps, k = 1..count.""" | |
| from transformers import TrainerCallback | |
| class EvenCheckpoints(TrainerCallback): | |
| def on_train_begin(self, args, state, control, **kwargs): | |
| # state.max_steps is final here; the last checkpoint always lands on the last step. | |
| self.steps = {max(1, round(state.max_steps * k / count)): k for k in range(1, count + 1)} | |
| def on_step_end(self, args, state, control, model=None, **kwargs): | |
| part = self.steps.get(state.global_step) | |
| if part is None: | |
| return control | |
| control.should_evaluate = True # Trainer evaluates right after this callback | |
| if state.is_world_process_zero: | |
| path = Path(output) / f"checkpoint-{part}of{count}-step{state.global_step}" | |
| save_adapter(model, tokenizer, recipe, path) | |
| print(f"Saved adapter checkpoint {path}", flush=True) | |
| return control | |
| return EvenCheckpoints() | |
| def prediction_logits(prediction): | |
| """Trainer may return logits alone or a tuple containing logits and auxiliary outputs.""" | |
| return prediction.predictions[0] if isinstance(prediction.predictions, tuple) else prediction.predictions | |
| def compute_metrics(prediction): | |
| return classification_metrics( | |
| prediction.label_ids.tolist(), prediction_logits(prediction).argmax(axis=-1).tolist() | |
| ) | |
| def softmax_rows(logits): | |
| import numpy as np | |
| shifted = logits - logits.max(axis=-1, keepdims=True) | |
| exp = np.exp(shifted) | |
| return exp / exp.sum(axis=-1, keepdims=True) | |
| def prediction_rows(rows, logits): | |
| """One record per input with the gold label, predicted label and class probabilities.""" | |
| probabilities = softmax_rows(logits.astype("float64")) | |
| keep = ("pair_id", "case_id", "domain", "policy_id", "conv_id", "synthetic", "selection_tier", | |
| "message", "policy", "label") | |
| return [ | |
| { | |
| **{key: row.get(key) for key in keep}, | |
| "predicted": LABELS[int(probs.argmax())], | |
| "probabilities": {name: float(probs[index]) for index, name in enumerate(LABELS)}, | |
| } | |
| for row, probs in zip(rows, probabilities, strict=True) | |
| ] | |
| def evaluation_report(records): | |
| """Metrics, confusion matrix and the human-readable sklearn report from prediction records.""" | |
| from sklearn.metrics import classification_report | |
| truth = [LABEL2ID[record["label"]] for record in records] | |
| predicted = [LABEL2ID[record["predicted"]] for record in records] | |
| return { | |
| "metrics": classification_metrics(truth, predicted), | |
| "confusion_matrix": confusion_matrix_counts(truth, predicted), | |
| "classification_report": classification_report( | |
| truth, predicted, labels=list(range(len(LABELS))), target_names=list(LABELS), | |
| digits=4, zero_division=0, | |
| ), | |
| } | |
| def run_training(config, model, tokenizer, base_config, collator, resume_from_checkpoint=None): | |
| """Train with HF Trainer and W&B logging; adapters saved at every checkpoint fraction.""" | |
| import torch | |
| import wandb | |
| from transformers import Trainer, TrainingArguments, set_seed | |
| from datasets import Dataset | |
| if not torch.cuda.is_available() or not torch.cuda.is_bf16_supported(): | |
| raise RuntimeError("Training requires a CUDA GPU with BF16 support (configured for A100)") | |
| splits = validate_inputs(config) | |
| set_seed(config["seed"]) | |
| def tokenize(batch): | |
| encoded = tokenize_pairs(tokenizer, batch["message"], batch["policy"], config["max_length"]) | |
| encoded["labels"] = batch["labels"] | |
| return encoded | |
| datasets = { | |
| name: Dataset.from_list(rows).map(tokenize, batched=True, remove_columns=list(rows[0])) | |
| for name, rows in splits.items() | |
| if rows | |
| } | |
| output = Path(config["output_dir"]) | |
| output.mkdir(parents=True, exist_ok=True) | |
| (output / "training_config.json").write_text(json.dumps(config, indent=2), encoding="utf-8") | |
| recipe = input_recipe(config, tokenizer, base_config) | |
| training_args = TrainingArguments( | |
| output_dir=str(output), | |
| run_name=config["run_name"], | |
| num_train_epochs=config["epochs"], | |
| learning_rate=config["learning_rate"], | |
| lr_scheduler_type=config["lr_scheduler_type"], | |
| warmup_ratio=config["warmup_ratio"], | |
| weight_decay=config["weight_decay"], | |
| max_grad_norm=config["max_grad_norm"], | |
| per_device_train_batch_size=config["batch_size"], | |
| per_device_eval_batch_size=config["eval_batch_size"], | |
| gradient_accumulation_steps=config["gradient_accumulation_steps"], | |
| bf16=True, | |
| tf32=True, | |
| gradient_checkpointing=config["gradient_checkpointing"], | |
| gradient_checkpointing_kwargs={"use_reentrant": False}, | |
| group_by_length=config.get("group_by_length", False), | |
| # Evaluation and saving are driven by the checkpoint callback, never by Trainer itself. | |
| eval_strategy="no", | |
| save_strategy="no", | |
| logging_steps=config["logging_steps"], | |
| report_to=["wandb"], | |
| seed=config["seed"], | |
| data_seed=config["seed"], | |
| optim="adamw_torch", | |
| dataloader_num_workers=config["dataloader_num_workers"], | |
| eval_accumulation_steps=8, | |
| ) | |
| trainer = Trainer( | |
| model=model, | |
| args=training_args, | |
| train_dataset=datasets["train"], | |
| eval_dataset=datasets["validation"], | |
| processing_class=tokenizer, | |
| data_collator=collator, | |
| compute_metrics=compute_metrics, | |
| callbacks=[quarter_checkpoint_callback(tokenizer, recipe, output, config["checkpoints_per_run"])], | |
| ) | |
| # Open the W&B run only after setup succeeds; the finally block closes it on errors too. | |
| run = wandb.init(project=config["wandb_project"], name=config["run_name"], config=config) | |
| (output / "wandb_run.json").write_text(json.dumps({"url": run.url, "id": run.id}), encoding="utf-8") | |
| try: | |
| torch.cuda.reset_peak_memory_stats() | |
| started = time.perf_counter() | |
| trained = trainer.train(resume_from_checkpoint=resume_from_checkpoint) | |
| torch.cuda.synchronize() | |
| hardware = { | |
| "gpu": torch.cuda.get_device_name(), | |
| "peak_allocated_gib": torch.cuda.max_memory_allocated() / 2**30, | |
| "peak_reserved_gib": torch.cuda.max_memory_reserved() / 2**30, | |
| "training_wall_seconds": time.perf_counter() - started, | |
| } | |
| (output / "hardware_metrics.json").write_text(json.dumps(hardware, indent=2), encoding="utf-8") | |
| wandb.log(hardware) | |
| trainer.save_metrics("train", trained.metrics) | |
| trainer.save_state() | |
| save_adapter(model, tokenizer, recipe, output / "final") # same weights as the last checkpoint | |
| for split in ("validation", "test_ood"): | |
| if split not in datasets: | |
| continue | |
| prediction = trainer.predict(datasets[split], metric_key_prefix=split) | |
| records = prediction_rows(splits[split], prediction_logits(prediction)) | |
| report = evaluation_report(records) | |
| write_rows(output / f"{split}_predictions.jsonl", records) | |
| (output / f"{split}_report.json").write_text(json.dumps(report, indent=2), encoding="utf-8") | |
| print(report["classification_report"], flush=True) | |
| trainer.log_metrics(split, prediction.metrics) | |
| wandb.log(prediction.metrics) | |
| wandb.log( | |
| { | |
| f"{split}_confusion_matrix": wandb.plot.confusion_matrix( | |
| y_true=[LABEL2ID[r["label"]] for r in records], | |
| preds=[LABEL2ID[r["predicted"]] for r in records], | |
| class_names=list(LABELS), | |
| ), | |
| f"{split}_errors": wandb.Table( | |
| columns=["message", "policy", "actual", "predicted", "p_violated"], | |
| data=[ | |
| [r["message"], r["policy"], r["label"], r["predicted"], | |
| r["probabilities"]["violated"]] | |
| for r in records if r["label"] != r["predicted"] | |
| ][:500], | |
| ), | |
| } | |
| ) | |
| finally: | |
| wandb.finish() | |
| def load_trained_classifier(checkpoint, device): | |
| """Rebuild base model + saved adapter/head from a checkpoint directory's input_format.json.""" | |
| import torch | |
| from peft import PeftModel | |
| from transformers import AutoModelForSequenceClassification, AutoTokenizer | |
| checkpoint = Path(checkpoint) | |
| recipe = json.loads((checkpoint / "input_format.json").read_text(encoding="utf-8")) | |
| tokenizer = AutoTokenizer.from_pretrained(checkpoint) | |
| dtype = torch.bfloat16 if device == "cuda" and torch.cuda.is_bf16_supported() else torch.float32 | |
| extra = {"reference_compile": False} if recipe["model_type"] == "modernbert" else {"use_cache": False} | |
| base = AutoModelForSequenceClassification.from_pretrained( | |
| recipe["model_name"], | |
| revision=recipe.get("model_revision"), | |
| num_labels=len(recipe["label2id"]), | |
| label2id=recipe["label2id"], | |
| id2label={index: label for label, index in recipe["label2id"].items()}, | |
| pad_token_id=tokenizer.pad_token_id, | |
| torch_dtype=dtype, | |
| attn_implementation=recipe.get("attention_implementation", "sdpa"), | |
| **extra, | |
| ) | |
| model = PeftModel.from_pretrained(base, checkpoint).to(device).eval() | |
| return model, tokenizer, recipe | |
| def predict_logits(model, tokenizer, rows, recipe, batch_size, device): | |
| """Batched inference in the saved input format; returns a numpy array of logits.""" | |
| import numpy as np | |
| import torch | |
| # Longest-first batches waste the least padding; results are put back in input order. | |
| order = sorted(range(len(rows)), key=lambda index: -len(rows[index]["message"]) - len(rows[index]["policy"])) | |
| logits = [None] * len(rows) | |
| with torch.inference_mode(): | |
| for start in range(0, len(order), batch_size): | |
| chunk = order[start : start + batch_size] | |
| batch = tokenize_pairs( | |
| tokenizer, | |
| [rows[index]["message"] for index in chunk], | |
| [rows[index]["policy"] for index in chunk], | |
| recipe["max_length"], | |
| recipe["instruction"], | |
| padding=True, | |
| return_tensors="pt", | |
| ).to(device) | |
| output = model(**batch).logits.float().cpu().numpy() | |
| for index, value in zip(chunk, output, strict=True): | |
| logits[index] = value | |
| print(f"predicted {min(start + batch_size, len(order))}/{len(order)}", flush=True) | |
| return np.stack(logits) | |
Xet Storage Details
- Size:
- 17.1 kB
- Xet hash:
- b624f4c76f9376393d3b601d6c52665e90a38e1235c03167ea0ba84741775b8d
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.