"""Main training script for the conspiracy language classifier.""" import os import random import time import numpy as np import torch from transformers import TrainingArguments, TrainerCallback, set_seed from config import CONFIG from data.preprocess import preprocess_pipeline from data.dataset import ConspiracyDataset, get_tokenizer from model.classifier import load_model, WeightedTrainer from model.utils import compute_metrics class TrainingLogger(TrainerCallback): """Prints a live metrics table for both training and validation each epoch.""" def __init__(self, train_dataset=None): self.epoch_train_losses = [] self.start_time = None self.header_printed = False self.train_dataset = train_dataset self.history = [] # list of dicts, one per epoch self._evaluating_train = False def on_train_begin(self, args, state, control, **kwargs): self.start_time = time.time() total_steps = state.max_steps total_epochs = args.num_train_epochs print(f"\nStarting {total_steps:,}-step training run " f"({int(total_epochs)} epochs)...\n") def on_log(self, args, state, control, logs=None, **kwargs): if logs is None: return # Accumulate training loss for epoch average if "loss" in logs: self.epoch_train_losses.append(logs["loss"]) # Print progress line periodically elapsed = time.time() - self.start_time elapsed_str = time.strftime("%H:%M:%S", time.gmtime(elapsed)) remaining = (elapsed / max(state.global_step, 1)) * (state.max_steps - state.global_step) remaining_str = time.strftime("%H:%M:%S", time.gmtime(remaining)) epoch_frac = state.global_step / (state.max_steps / args.num_train_epochs) speed = state.global_step / max(elapsed, 1) print( f"\r [{state.global_step}/{state.max_steps} " f"{elapsed_str} < {remaining_str}, " f"{speed:.2f} it/s, " f"Epoch {epoch_frac:.1f}/{int(args.num_train_epochs)}] " f"loss: {logs['loss']:.4f}", end="", flush=True, ) def on_evaluate(self, args, state, control, metrics=None, model=None, **kwargs): if metrics is None: return # Guard: skip if this is our own train-set evaluation (prevents infinite loop) if self._evaluating_train: return # Evaluate on a small sample of training set (for overfitting diagnostics) train_metrics = {} if self.train_dataset is not None: if hasattr(self, "_trainer") and self._trainer is not None: self._evaluating_train = True try: sample_size = min(1000, len(self.train_dataset)) train_sample = torch.utils.data.Subset( self.train_dataset, list(range(sample_size)), ) train_metrics = self._trainer.evaluate( eval_dataset=train_sample, metric_key_prefix="train" ) finally: self._evaluating_train = False # Print header once if not self.header_printed: print() # newline after progress bar print() print(f" {'':>7} {'──── Training ────':^36s} {'──── Validation ────':^36s}") print(f" {'Epoch':<7} {'Loss':<10}{'Acc':<9}{'F1':<9}{'Prec':<9}" f" {'Loss':<10}{'Acc':<9}{'F1':<9}{'Prec':<9}{'Recall':<9}{'AUC':<9}") print(f" {'─' * 99}") self.header_printed = True # Average training loss from logged steps avg_train_loss = ( sum(self.epoch_train_losses) / len(self.epoch_train_losses) if self.epoch_train_losses else 0.0 ) self.epoch_train_losses = [] epoch = int(state.epoch) if state.epoch else "?" # Training metrics (from evaluate on train set) t_loss = train_metrics.get("train_loss", avg_train_loss) t_acc = train_metrics.get("train_accuracy", 0.0) t_f1 = train_metrics.get("train_f1", 0.0) t_prec = train_metrics.get("train_precision", 0.0) # Validation metrics v_loss = metrics.get("eval_loss", 0.0) v_acc = metrics.get("eval_accuracy", 0.0) v_f1 = metrics.get("eval_f1", 0.0) v_prec = metrics.get("eval_precision", 0.0) v_rec = metrics.get("eval_recall", 0.0) v_auc = metrics.get("eval_roc_auc", 0.0) # Save history for summary self.history.append({ "epoch": epoch, "train_loss": t_loss, "train_acc": t_acc, "train_f1": t_f1, "val_loss": v_loss, "val_acc": v_acc, "val_f1": v_f1, "val_prec": v_prec, "val_rec": v_rec, "val_auc": v_auc, }) # Format training metrics — show dashes if not available if train_metrics: t_str = (f"{t_loss:<10.4f}{t_acc:<9.4f}{t_f1:<9.4f}{t_prec:<9.4f}") else: t_str = (f"{avg_train_loss:<10.4f}{'—':<9}{'—':<9}{'—':<9}") print(f" {epoch:<7} {t_str}" f" {v_loss:<10.4f}{v_acc:<9.4f}{v_f1:<9.4f}{v_prec:<9.4f}{v_rec:<9.4f}{v_auc:<9.4f}") def on_train_end(self, args, state, control, **kwargs): elapsed = time.time() - self.start_time elapsed_str = time.strftime("%H:%M:%S", time.gmtime(elapsed)) # ── Summary ─────────────────────────────────────────────────── print(f"\n Training finished in {elapsed_str}") print(f" Best metric ({args.metric_for_best_model}): " f"{state.best_metric:.4f} at step {state.best_model_checkpoint}") if len(self.history) >= 2: print(f"\n ── Overfitting Check ──") best_val = max(self.history, key=lambda h: h["val_f1"]) last = self.history[-1] print(f" Best val F1: {best_val['val_f1']:.4f} (epoch {best_val['epoch']})") print(f" Final val F1: {last['val_f1']:.4f} (epoch {last['epoch']})") if best_val["train_acc"] > 0 and best_val["val_acc"] > 0: gap = best_val["train_acc"] - best_val["val_acc"] print(f" Train-Val accuracy gap at best epoch: {gap:+.4f}") if gap > 0.05: print(f" ⚠ Gap > 5% — model may still be overfitting") else: print(f" ✓ Gap looks healthy") def resolve_mixed_precision(): """Auto-detect best mixed precision mode for the current GPU.""" if not torch.cuda.is_available(): print(" Mixed precision: disabled (no GPU)") return False, False capability = torch.cuda.get_device_capability() gpu_name = torch.cuda.get_device_name() print(f" GPU: {gpu_name} (compute capability {capability[0]}.{capability[1]})") # bf16 needs compute capability >= 8.0 (A100, A10, etc.) if capability[0] >= 8: print(" Mixed precision: bf16") return False, True else: # T4 (7.5) and older: fp16 can be unstable with DeBERTa-v3. # Use fp32 for safety — slower but correct. print(" Mixed precision: fp32 (DeBERTa-v3 unstable with fp16 on this GPU)") return False, False def main(): # Full reproducibility — seed every RNG that touches training seed = CONFIG["seed"] random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed_all(seed) set_seed(seed) # transformers/HF Trainer torch.backends.cudnn.deterministic = True torch.backends.cudnn.benchmark = False # ── Phase 1: Data Preprocessing ─────────────────────────────── print("=" * 60) print(" Phase 1: Data Preprocessing") print("=" * 60) train_df, val_df, test_df, covid_test_df, class_weights = preprocess_pipeline() # Save splits for later evaluation os.makedirs(CONFIG["output_dir"], exist_ok=True) test_df.to_csv(os.path.join(CONFIG["output_dir"], "test_split.csv"), index=False) covid_test_df.to_csv(os.path.join(CONFIG["output_dir"], "covid_test_split.csv"), index=False) # ── Phase 2: Preparing Model and Data ───────────────────────── print("\n" + "=" * 60) print(" Phase 2: Preparing Model and Data") print("=" * 60) tokenizer = get_tokenizer() train_dataset = ConspiracyDataset(train_df, tokenizer, apply_debiasing=True) val_dataset = ConspiracyDataset(val_df, tokenizer) print(f" Train samples: {len(train_dataset):,}") print(f" Val samples: {len(val_dataset):,}") print(f" Class weights: [{class_weights[0]:.4f}, {class_weights[1]:.4f}]") # Show class balance train_label_counts = train_df["label"].value_counts().sort_index() print(f" Train label 0 (non-conspiracy): {train_label_counts.get(0, 0):,}") print(f" Train label 1 (conspiracy): {train_label_counts.get(1, 0):,}") # Show source distribution print(f"\n Source distribution in training set:") for src, count in train_df["source"].value_counts().items(): print(f" {src:<25} {count:>6,}") # 3. Model print(f"\n Loading model: {CONFIG['model_name']}...") model = load_model() total_params = sum(p.numel() for p in model.parameters()) trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad) print(f" Total parameters: {total_params:,}") print(f" Trainable parameters: {trainable_params:,}") # Compute training info steps_per_epoch = len(train_dataset) // ( CONFIG["per_device_train_batch_size"] * CONFIG["gradient_accumulation_steps"] ) total_steps = steps_per_epoch * CONFIG["num_epochs"] print(f"\n Steps per epoch: {steps_per_epoch:,}") print(f" Total steps: {total_steps:,}") print(f" Effective batch: {CONFIG['per_device_train_batch_size'] * CONFIG['gradient_accumulation_steps']}") print(f" Learning rate: {CONFIG['learning_rate']}") print(f" Epochs: {CONFIG['num_epochs']}") # ── Phase 3: Training ───────────────────────────────────────── use_fp16, use_bf16 = resolve_mixed_precision() training_args = TrainingArguments( output_dir=CONFIG["output_dir"], num_train_epochs=CONFIG["num_epochs"], per_device_train_batch_size=CONFIG["per_device_train_batch_size"], per_device_eval_batch_size=CONFIG["per_device_train_batch_size"] * 2, gradient_accumulation_steps=CONFIG["gradient_accumulation_steps"], learning_rate=CONFIG["learning_rate"], weight_decay=CONFIG["weight_decay"], warmup_ratio=CONFIG["warmup_ratio"], max_grad_norm=CONFIG.get("max_grad_norm", 1.0), fp16=use_fp16, bf16=use_bf16, eval_strategy="epoch", save_strategy="epoch", load_best_model_at_end=True, metric_for_best_model="f1", greater_is_better=True, save_total_limit=3, logging_steps=20, report_to="none", seed=CONFIG["seed"], dataloader_num_workers=2, disable_tqdm=True, # we use our own logger ) print("\n" + "=" * 60) print(" Phase 3: Training") print("=" * 60) training_logger = TrainingLogger(train_dataset=train_dataset) trainer = WeightedTrainer( class_weights=class_weights, model=model, args=training_args, train_dataset=train_dataset, eval_dataset=val_dataset, compute_metrics=compute_metrics, callbacks=[training_logger], ) # Give the logger a reference to the trainer so it can evaluate on train set training_logger._trainer = trainer trainer.train() # ── Phase 4: Save ───────────────────────────────────────────── print("\n" + "=" * 60) print(" Phase 4: Saving Model") print("=" * 60) best_model_path = os.path.join(CONFIG["output_dir"], "best_model") trainer.save_model(best_model_path) tokenizer.save_pretrained(best_model_path) print(f" Best model saved to: {best_model_path}") # Show saved size (LoRA adapter should be ~10-15 MB) total_size = sum( os.path.getsize(os.path.join(best_model_path, f)) for f in os.listdir(best_model_path) if os.path.isfile(os.path.join(best_model_path, f)) ) print(f" Saved model size: {total_size / 1024 / 1024:.1f} MB") # ── Phase 5: Validation Results ─────────────────────────────── print("\n" + "=" * 60) print(" Phase 5: Final Validation") print("=" * 60) val_results = trainer.evaluate() for key, value in sorted(val_results.items()): if isinstance(value, float): print(f" {key:<25} {value:.4f}") else: print(f" {key:<25} {value}") print("\n\n") print("&" * 80) print("&" * 80) print("&&" + " " * 76 + "&&") print("&&" + "EVALUATION".center(76) + "&&") print("&&" + " " * 76 + "&&") print("&" * 80) print("&" * 80) print("\n") # ── Phase 6: Full Evaluation ────────────────────────────────── from evaluate import main as evaluate_main evaluate_main() if __name__ == "__main__": main()