train.py

#!/usr/bin/env python3
"""Gotcha Clause Extractor: Comprehensive Production Training Pipeline

Fine-tunes transformer models (ELECTRA-Small, TinyBERT, BERT-Mini, BERT-Tiny)
for legal sequence labeling (BIO tag token classification) with:
1. Unified Corpus: CodeHima/TOS_Dataset + EE21/ToS-Summaries + Local Synthetic Data.
2. Symmetric context mixing (prevents positional inductive bias).
3. Exact multi-instance sub-word token BIO alignment.
4. Inverse-sqrt class-weighted CrossEntropy loss ([0.38, 4.70, 1.00]).
5. Optuna hyperparameter optimization with complete trial telemetry logging.
6. Comprehensive JSON metric logging (trials, learning curves, hold-out test metrics).
7. Automatic checkpoint cleanup to maintain lean disk footprint.
"""

import os
import re
import sys
import json
import copy
import random
import shutil
import argparse
import datetime
import time
try:
    import psutil
except ImportError:
    psutil = None
from typing import Dict, Any, List, Optional

import numpy as np
import pandas as pd
import torch
import torch.nn as nn
from datasets import Dataset, concatenate_datasets
import nltk
from nltk.tokenize import sent_tokenize
try:
    nltk.download('punkt', quiet=True)
    nltk.download('punkt_tab', quiet=True)
except Exception:
    pass
import evaluate
import optuna
from transformers import (
    AutoTokenizer,
    AutoModelForTokenClassification,
    BertForTokenClassification,
    TrainingArguments,
    Trainer,
    DataCollatorForTokenClassification,
    TrainerCallback,
    PrinterCallback,
    logging as tf_logging
)

from gotcha.constants import label2id, id2label, AVAILABLE_MODELS
from gotcha.preprocessor import clean_text_pipeline

os.environ["HF_HUB_DISABLE_PROGRESS_BARS"] = "1"
os.environ["HF_HUB_DISABLE_SYMLINKS_WARNING"] = "1"
tf_logging.set_verbosity_error()
tf_logging.disable_progress_bar()
optuna.logging.set_verbosity(optuna.logging.WARNING)

TOKENIZER_MAP = {
    "prajjwal1/bert-tiny": "google-bert/bert-base-uncased",
    "prajjwal1/bert-mini": "google-bert/bert-base-uncased",
    "bert-tiny": "google-bert/bert-base-uncased",
    "bert-mini": "google-bert/bert-base-uncased",
}

MODEL_CLASS_MAP = {
    "prajjwal1/bert-tiny": BertForTokenClassification,
    "prajjwal1/bert-mini": BertForTokenClassification,
    "bert-tiny": BertForTokenClassification,
    "bert-mini": BertForTokenClassification,
}

MODEL_ARCH_MAP = {
    "google/electra-small-discriminator": {"backbone_attr": "electra", "max_layers": 12},
    "electra-small":                     {"backbone_attr": "electra", "max_layers": 12},
    "prajjwal1/bert-tiny":                {"backbone_attr": "bert",    "max_layers": 2},
    "bert-tiny":                          {"backbone_attr": "bert",    "max_layers": 2},
    "prajjwal1/bert-mini":                {"backbone_attr": "bert",    "max_layers": 4},
    "bert-mini":                          {"backbone_attr": "bert",    "max_layers": 4},
    "huawei-noah/TinyBERT_General_4L_312D": {"backbone_attr": "bert",  "max_layers": 4},
    "tinybert":                           {"backbone_attr": "bert",    "max_layers": 4},
}

MODEL_PRETRAINED_MAP = {
    "electra-small": "google/electra-small-discriminator",
    "tinybert": "huawei-noah/TinyBERT_General_4L_312D",
    "bert-mini": "prajjwal1/bert-mini",
    "bert-tiny": "prajjwal1/bert-tiny",
}

STOPWORDS = set('a an the in on at of and or is are was were to for with by as this that it be from have has had not but'.split())

seqeval_metric = evaluate.load("seqeval")


def build_unified_dataset(synthetic_path="corpus/synthetic_gotchas.json", seed=42):
    """Build and partition a unified legal dataset combining CodeHima, EE21, and Synthetic data."""
    print("=" * 65)
    print("  BUILDING UNIFIED TRAINING CORPUS")
    print("=" * 65)

    standardized_data = []
    neutral_sentences = []
    gotcha_sentences = []

    # 1. Ingest CodeHima
    print("Fetching CodeHima/TOS_Dataset from Hugging Face...")
    try:
        splits = {'train': 'data/train-00000-of-00001.parquet'}
        df_codehima = pd.read_parquet("hf://datasets/CodeHima/TOS_Dataset/" + splits["train"])

        for _, row in df_codehima.iterrows():
            sentence = clean_text_pipeline(str(row['sentence']))
            if not sentence or len(sentence) < 10:
                continue
            if row['unfairness_level'] in ['potentially_unfair', 'clearly_unfair']:
                gotcha_sentences.append(sentence)
            else:
                neutral_sentences.append(sentence)
                standardized_data.append({"text": sentence, "gotchas": []})

        # Symmetric context wrapping to eliminate positional prior
        random.seed(seed)
        for gotcha_sent in gotcha_sentences:
            r = random.random()
            if r < 0.40:
                c = random.choice(neutral_sentences)
                comb = c + " " + gotcha_sent
            elif r < 0.80:
                c = random.choice(neutral_sentences)
                comb = gotcha_sent + " " + c
            else:
                c1 = random.choice(neutral_sentences)
                c2 = random.choice(neutral_sentences)
                comb = c1 + " " + gotcha_sent + " " + c2
            standardized_data.append({"text": comb, "gotchas": [gotcha_sent]})

        print(f"  CodeHima parsed: {len(neutral_sentences)} neutral, {len(gotcha_sentences)} gotchas (symmetrically wrapped)")
    except Exception as e:
        print(f"  Warning loading CodeHima: {e}")

    # 2. Ingest EE21/ToS-Summaries
    print("Fetching EE21/ToS-Summaries from Hugging Face...")
    ee21_gotchas = 0
    ee21_neutrals = 0
    try:
        df_summaries = pd.read_json("hf://datasets/EE21/ToS-Summaries/dataset.json", lines=True)
        for _, row in df_summaries.iterrows():
            text = clean_text_pipeline(str(row['plain_text']))
            summary = str(row['summary']).lower()
            sum_words = set(re.findall(r'\b[a-z]{3,}\b', summary)) - STOPWORDS
            if not sum_words or len(text) < 30:
                continue

            sentences = sent_tokenize(text)
            for s in sentences:
                if len(s) < 25 or len(s) > 500:
                    continue
                s_words = set(re.findall(r'\b[a-z]{3,}\b', s.lower())) - STOPWORDS
                overlap = len(sum_words & s_words)
                if overlap >= 4:
                    standardized_data.append({"text": s, "gotchas": [s]})
                    ee21_gotchas += 1
                elif overlap == 0:
                    standardized_data.append({"text": s, "gotchas": []})
                    ee21_neutrals += 1

        print(f"  EE21 parsed: {ee21_gotchas} gotcha spans, {ee21_neutrals} neutral segments")
    except Exception as e:
        print(f"  Warning loading EE21: {e}")

    print(f"  Total human-labeled samples: {len(standardized_data)}")
    raw_original_dataset = Dataset.from_list(standardized_data)
    dataset_splits = raw_original_dataset.train_test_split(test_size=0.15, seed=seed)

    # 3. Augment Train Split Exclusively with Synthetic Gotchas & Legal Prefix Noise
    if os.path.exists(synthetic_path):
        print(f"Merging synthetic gotchas from {synthetic_path} strictly into train split...")
        try:
            with open(synthetic_path, "r", encoding="utf-8") as f:
                synthetic_data = json.load(f)

            # Prefix noise augmentation: prepend numbering / section headers to a subset of samples
            # Teaches the token classifier that leading legal numbering prefixes are strictly 'O'
            random.seed(seed)
            legal_prefixes = [
                "1. ", "2. ", "12. ", "14.1 ", "Section 4. ", "Section 12. ",
                "Article III: ", "Clause 2 - ", "(a) ", "(b) ", "[i] ", "[ii] "
            ]
            augmented_samples = []
            for item in synthetic_data:
                augmented_samples.append(item)
                if random.random() < 0.15:
                    pfx = random.choice(legal_prefixes)
                    augmented_samples.append({
                        "text": pfx + item["text"],
                        "gotchas": item.get("gotchas", [])
                    })

            synthetic_dataset = Dataset.from_list(augmented_samples)
            dataset_splits["train"] = concatenate_datasets([dataset_splits["train"], synthetic_dataset])
            print(f"  Added {len(augmented_samples)} synthetic & prefix-augmented samples. Train split size: {len(dataset_splits['train'])}")
        except Exception as e:
            print(f"  Warning loading synthetic data: {e}")

    print(f"Final partitioned dataset: train={len(dataset_splits['train'])}, test (100% human)={len(dataset_splits['test'])}")
    return dataset_splits


def tokenize_and_align_labels(example, tokenizer, max_length=512):
    """Align multi-occurrence gotchas to sub-word token sequences."""
    text = example['text']
    gotchas = example.get('gotchas', [])

    tokenized_inputs = tokenizer(
        text,
        truncation=True,
        max_length=max_length,
        return_offsets_mapping=True
    )

    offsets = tokenized_inputs["offset_mapping"]
    labels = [label2id['O']] * len(offsets)

    for gotcha in gotchas:
        if not gotcha:
            continue
        gotcha_len = len(gotcha)
        pos = 0
        while True:
            start_char = text.find(gotcha, pos)
            if start_char == -1:
                break
            end_char = start_char + gotcha_len

            gotcha_started = False
            for idx, (start, end) in enumerate(offsets):
                if start == end:
                    labels[idx] = -100
                    continue
                if start >= start_char and start < end_char:
                    if not gotcha_started:
                        labels[idx] = label2id['B-RISK']
                        gotcha_started = True
                    else:
                        labels[idx] = label2id['I-RISK']

            pos = end_char

    tokenized_inputs["labels"] = labels
    del tokenized_inputs["offset_mapping"]
    return tokenized_inputs


class WeightedLossTrainer(Trainer):
    def __init__(self, *args, class_weights=None, **kwargs):
        super().__init__(*args, **kwargs)
        if class_weights is None:
            self._class_weights = torch.tensor([0.38, 4.70, 1.00])
        else:
            self._class_weights = torch.tensor(class_weights)
        self._cached_loss_fct = None

    def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
        labels = inputs.get("labels")
        outputs = model(**inputs)
        logits = outputs.get("logits")
        device = logits.device

        if self._cached_loss_fct is None or self._cached_loss_fct.weight.device != device:
            self._cached_loss_fct = nn.CrossEntropyLoss(
                weight=self._class_weights.to(device)
            )

        if labels is not None:
            active_loss = inputs.get("attention_mask").view(-1) == 1
            active_logits = logits.view(-1, self.model.config.num_labels)
            active_labels = torch.where(
                active_loss, labels.view(-1),
                torch.tensor(-100, device=device)
            )
            loss = self._cached_loss_fct(active_logits, active_labels)
        else:
            loss = outputs.loss

        return (loss, outputs) if return_outputs else loss


class ComprehensiveMetricsCallback(TrainerCallback):
    """Collects and displays rich training telemetry:
    - F1, Precision, Recall, Accuracy, Train & Eval Loss, Learning Rate
    - Hardware Telemetry: Live & Peak VRAM allocated/reserved (MB/GB), GPU utilization %, Process & System RAM
    - Performance Telemetry: Throughput (samples/sec), epoch duration, ETA
    """
    def __init__(self, mode="final", total_train_samples=0, num_epochs=10):
        self.mode = mode
        self.total_train_samples = total_train_samples
        self.num_epochs = num_epochs
        self.history = {}
        self.trials_history = {}
        self.current_trial = 0
        self.current_trial_epochs = {}
        self.train_start_time = None
        self.epoch_start_time = None
        self.epoch_durations = []
        self.peak_vram_allocated_mb = 0.0
        self.peak_vram_reserved_mb = 0.0
        self.peak_process_ram_mb = 0.0

    def on_train_begin(self, args, state, control, **kwargs):
        self.train_start_time = time.time()
        self.epoch_start_time = time.time()
        if torch.cuda.is_available():
            torch.cuda.reset_peak_memory_stats()

    def on_epoch_begin(self, args, state, control, **kwargs):
        self.epoch_start_time = time.time()

    def on_log(self, args, state, control, logs=None, **kwargs):
        if not logs:
            return
        ep = int(round(state.epoch)) if state.epoch is not None else 0
        if ep <= 0:
            return

        target_dict = self.current_trial_epochs if self.mode == "trial" else self.history
        if ep not in target_dict:
            target_dict[ep] = {"Epoch": ep}

        entry = target_dict[ep]
        if "loss" in logs:
            entry["Train_Loss"] = round(float(logs["loss"]), 4)
        if "eval_loss" in logs:
            entry["Eval_Loss"] = round(float(logs["eval_loss"]), 4)
        if "eval_f1" in logs:
            entry["F1"] = round(float(logs["eval_f1"]), 4)
        if "eval_precision" in logs:
            entry["Precision"] = round(float(logs["eval_precision"]), 4)
        if "eval_recall" in logs:
            entry["Recall"] = round(float(logs["eval_recall"]), 4)
        if "eval_accuracy" in logs:
            entry["Accuracy"] = round(float(logs["eval_accuracy"]), 4)
        if "learning_rate" in logs:
            entry["Learning_Rate"] = round(float(logs["learning_rate"]), 7)

        now = time.time()
        epoch_dur = max(0.1, now - (self.epoch_start_time or now))
        entry["Epoch_Duration_Sec"] = round(epoch_dur, 2)
        if self.total_train_samples > 0:
            entry["Samples_Per_Sec"] = round(self.total_train_samples / epoch_dur, 1)

        # Hardware & VRAM Telemetry
        if torch.cuda.is_available():
            alloc_mb = torch.cuda.memory_allocated() / (1024 ** 2)
            max_alloc_mb = torch.cuda.max_memory_allocated() / (1024 ** 2)
            res_mb = torch.cuda.memory_reserved() / (1024 ** 2)
            max_res_mb = torch.cuda.max_memory_reserved() / (1024 ** 2)
            total_vram_gb = torch.cuda.get_device_properties(0).total_memory / (1024 ** 3)
            entry["VRAM_Allocated_MB"] = round(alloc_mb, 1)
            entry["VRAM_Peak_Allocated_MB"] = round(max_alloc_mb, 1)
            entry["VRAM_Reserved_MB"] = round(res_mb, 1)
            entry["VRAM_Peak_Reserved_MB"] = round(max_res_mb, 1)
            entry["VRAM_Peak_GB"] = round(max_res_mb / 1024, 2)
            entry["VRAM_Util_Pct"] = round((max_res_mb / (total_vram_gb * 1024)) * 100, 1)
            self.peak_vram_allocated_mb = max(self.peak_vram_allocated_mb, max_alloc_mb)
            self.peak_vram_reserved_mb = max(self.peak_vram_reserved_mb, max_res_mb)

        if psutil is not None:
            try:
                proc_ram_mb = psutil.Process().memory_info().rss / (1024 ** 2)
                entry["Process_RAM_MB"] = round(proc_ram_mb, 1)
                self.peak_process_ram_mb = max(self.peak_process_ram_mb, proc_ram_mb)
            except Exception:
                pass

        if "eval_f1" in logs:
            total_epochs = int(args.num_train_epochs) if args.num_train_epochs else self.num_epochs
            self.epoch_durations.append(epoch_dur)
            avg_dur = sum(self.epoch_durations) / len(self.epoch_durations)
            remaining_eps = max(0, total_epochs - ep)
            eta_sec = int(remaining_eps * avg_dur)
            eta_str = f"{eta_sec // 60}m {eta_sec % 60:02d}s" if eta_sec >= 60 else f"{eta_sec}s"

            prefix = f"  [Trial {self.current_trial + 1} | Ep {ep}/{total_epochs}]" if self.mode == "trial" else f"  [Epoch {ep}/{total_epochs}]"
            train_l = entry.get("Train_Loss", "—")
            eval_l = entry.get("Eval_Loss", "—")
            f1 = entry.get("F1", 0.0)
            prec = entry.get("Precision", 0.0)
            rec = entry.get("Recall", 0.0)
            samples_s = entry.get("Samples_Per_Sec", "—")

            if torch.cuda.is_available():
                peak_gb = entry.get("VRAM_Peak_GB", 0.0)
                tot_gb = total_vram_gb
                util_pct = entry.get("VRAM_Util_Pct", 0.0)
                vram_info = f"VRAM: {entry.get('VRAM_Allocated_MB', 0):,.0f} MB alloc / {entry.get('VRAM_Peak_Reserved_MB', 0):,.0f} MB peak ({peak_gb:.2f} GB / {tot_gb:.2f} GB, {util_pct:.1f}%)"
            else:
                proc_ram = entry.get("Process_RAM_MB", 0.0)
                vram_info = f"RAM: {proc_ram:,.0f} MB process"

            print(f"{prefix} Train: {train_l} | Eval: {eval_l} | F1: {f1:.4f} | Prec: {prec:.4f} | Rec: {rec:.4f}")
            print(f"      {vram_info} | {samples_s} samples/s | Ep: {epoch_dur:.1f}s | ETA: {eta_str}")
            self.epoch_start_time = time.time()

    def on_train_end(self, args, state, control, **kwargs):
        if self.mode == "trial":
            if self.current_trial_epochs:
                sorted_epochs = [self.current_trial_epochs[k] for k in sorted(self.current_trial_epochs.keys())]
                self.trials_history[self.current_trial] = {
                    "epochs": sorted_epochs,
                    "peak_vram_mb": round(self.peak_vram_reserved_mb, 1),
                    "peak_vram_gb": round(self.peak_vram_reserved_mb / 1024, 2),
                    "peak_allocated_mb": round(self.peak_vram_allocated_mb, 1),
                    "trial_duration_sec": round(time.time() - (self.train_start_time or time.time()), 1)
                }
                self.current_trial += 1
                self.current_trial_epochs = {}
                self.epoch_durations = []
                self.peak_vram_allocated_mb = 0.0
                self.peak_vram_reserved_mb = 0.0

    def get_final_history_list(self):
        return [self.history[k] for k in sorted(self.history.keys())]

    def reset(self):
        self.history = {}
        self.trials_history = {}
        self.current_trial = 0
        self.current_trial_epochs = {}
        self.epoch_durations = []
        self.peak_vram_allocated_mb = 0.0
        self.peak_vram_reserved_mb = 0.0


def compute_metrics(p):
    predictions, labels = p
    predictions = np.argmax(predictions, axis=2)

    true_predictions = [
        [id2label[p_idx] for (p_idx, l) in zip(prediction, label) if l != -100]
        for prediction, label in zip(predictions, labels)
    ]
    true_labels = [
        [id2label[l] for (p_idx, l) in zip(prediction, label) if l != -100]
        for prediction, label in zip(predictions, labels)
    ]
    results = seqeval_metric.compute(predictions=true_predictions, references=true_labels)
    return {
        "precision": float(results.get("overall_precision", 0.0)),
        "recall": float(results.get("overall_recall", 0.0)),
        "f1": float(results.get("overall_f1", 0.0)),
        "accuracy": float(results.get("overall_accuracy", 0.0)),
    }


def freeze_layers(model, model_checkpoint, num_frozen_layers):
    """Dynamically freeze embeddings + specified number of bottom encoder layers."""
    arch = MODEL_ARCH_MAP.get(model_checkpoint)
    if arch is None:
        return

    backbone = getattr(model, arch["backbone_attr"], None)
    if backbone is None:
        return

    if hasattr(backbone, "embeddings"):
        for param in backbone.embeddings.parameters():
            param.requires_grad = False

    if hasattr(backbone, "encoder") and hasattr(backbone.encoder, "layer"):
        actual_layers = len(backbone.encoder.layer)
        num_to_freeze = min(num_frozen_layers, actual_layers)
        for i in range(num_to_freeze):
            for param in backbone.encoder.layer[i].parameters():
                param.requires_grad = False


def evaluate_model_on_test(model_key: str, dataset_splits, output_base_dir: str = "./gotcha-extractor-model", device: str = "cuda"):
    """Run hold-out test evaluation on an already trained model and update metrics JSON."""
    model_dir = os.path.join(output_base_dir, model_key)
    if not os.path.exists(model_dir):
        print(f"[SKIP] Model directory {model_dir} does not exist.")
        return None

    print(f"\n--- Evaluating {model_key.upper()} on 100% Human Hold-Out Test Split ---")
    pretrained_id = MODEL_PRETRAINED_MAP.get(model_key, model_key)
    tok_name = TOKENIZER_MAP.get(pretrained_id, pretrained_id)
    tokenizer = AutoTokenizer.from_pretrained(tok_name)
    data_collator = DataCollatorForTokenClassification(tokenizer=tokenizer)

    tokenized_test = dataset_splits["test"].map(
        lambda x: tokenize_and_align_labels(x, tokenizer),
        remove_columns=dataset_splits["test"].column_names
    )

    model_cls = MODEL_CLASS_MAP.get(pretrained_id, AutoModelForTokenClassification)
    model = model_cls.from_pretrained(model_dir).to(device)

    eval_args = TrainingArguments(
        output_dir=os.path.join(output_base_dir, f"{model_key}-eval-temp"),
        per_device_eval_batch_size=8,
        report_to="none",
    )

    eval_trainer = WeightedLossTrainer(
        model=model,
        args=eval_args,
        eval_dataset=tokenized_test,
        processing_class=tokenizer,
        data_collator=data_collator,
        compute_metrics=compute_metrics,
    )
    eval_trainer.remove_callback(PrinterCallback)

    test_results = eval_trainer.evaluate(eval_dataset=tokenized_test)
    if os.path.exists(eval_args.output_dir):
        shutil.rmtree(eval_args.output_dir)

    test_eval = {
        "f1": round(float(test_results.get("eval_f1", 0.0)), 4),
        "precision": round(float(test_results.get("eval_precision", 0.0)), 4),
        "recall": round(float(test_results.get("eval_recall", 0.0)), 4),
        "loss": round(float(test_results.get("eval_loss", 0.0)), 4),
        "accuracy": round(float(test_results.get("eval_accuracy", 0.0)), 4),
    }
    print(f"  Hold-Out Test Results: F1={test_eval['f1']} | Recall={test_eval['recall']} | Precision={test_eval['precision']} | Loss={test_eval['loss']}")

    # Update existing metrics JSON
    metrics_file = os.path.join(output_base_dir, f"{model_key}_metrics.json")
    metrics_data = {}
    if os.path.exists(metrics_file):
        try:
            with open(metrics_file, "r", encoding="utf-8") as f:
                metrics_data = json.load(f)
        except Exception:
            metrics_data = {}

    metrics_data["test_eval"] = test_eval
    summary = metrics_data.get("summary", {})
    summary["test_f1"] = test_eval["f1"]
    summary["test_recall"] = test_eval["recall"]
    summary["test_precision"] = test_eval["precision"]
    metrics_data["summary"] = summary

    with open(metrics_file, "w", encoding="utf-8") as f:
        json.dump(metrics_data, f, indent=2)

    return test_eval


def train_single_model(
    model_key: str,
    dataset_splits,
    output_base_dir: str = "./gotcha-extractor-model",
    epochs: int = 10,
    batch_size: int = 8,
    learning_rate: float = 5e-5,
    weight_decay: float = 0.01,
    warmup_ratio: float = 0.15,
    num_frozen_layers: int = 0,
    use_optuna: bool = False,
    n_trials: int = 10,
    device: str = "cuda"
):
    """Train a designated model with comprehensive Optuna HPO and complete metric recording."""
    pretrained_id = MODEL_PRETRAINED_MAP.get(model_key, model_key)
    print("\n" + "=" * 65)
    print(f"  COMMENCING TRAINING: {model_key.upper()} ({pretrained_id})")
    print(f"  Epochs: {epochs} | Batch Size: {batch_size} | Default LR: {learning_rate} | Optuna: {use_optuna}")
    print("=" * 65)

    tok_name = TOKENIZER_MAP.get(pretrained_id, pretrained_id)
    model_cls = MODEL_CLASS_MAP.get(pretrained_id, AutoModelForTokenClassification)

    tokenizer = AutoTokenizer.from_pretrained(tok_name)
    data_collator = DataCollatorForTokenClassification(tokenizer=tokenizer)

    print("Tokenizing train and test splits...")
    tokenized_train = dataset_splits["train"].map(
        lambda x: tokenize_and_align_labels(x, tokenizer),
        remove_columns=dataset_splits["train"].column_names
    )
    tokenized_test = dataset_splits["test"].map(
        lambda x: tokenize_and_align_labels(x, tokenizer),
        remove_columns=dataset_splits["test"].column_names
    )

    arch = MODEL_ARCH_MAP.get(model_key, {})
    max_freezable = arch.get("max_layers", 0)
    output_dir = os.path.join(output_base_dir, model_key)
    os.makedirs(output_dir, exist_ok=True)

    # Read existing metrics JSON if it exists to preserve prior trials when running without Optuna
    metrics_file = os.path.join(output_base_dir, f"{model_key}_metrics.json")
    existing_metrics = {}
    if os.path.exists(metrics_file):
        try:
            with open(metrics_file, "r", encoding="utf-8") as f:
                existing_metrics = json.load(f)
        except Exception:
            existing_metrics = {}

    best_f1 = 0.0
    best_hp = {
        "learning_rate": learning_rate,
        "weight_decay": weight_decay,
        "per_device_train_batch_size": batch_size,
        "warmup_ratio": warmup_ratio,
        "num_frozen_layers": num_frozen_layers,
    }

    trials_dict = existing_metrics.get("trials", {})

    if use_optuna:
        print(f"\n--- Running Optuna Hyperparameter Search ({n_trials} trials, {epochs} epochs each) ---")
        search_dir = os.path.join(output_base_dir, f"{model_key}-optuna-search")
        trial_metrics_cb = ComprehensiveMetricsCallback(
            mode="trial",
            total_train_samples=len(tokenized_train),
            num_epochs=epochs
        )

        trial_records = {}

        def model_init(trial=None):
            m = model_cls.from_pretrained(
                pretrained_id,
                num_labels=len(label2id),
                id2label=id2label,
                label2id=label2id,
                ignore_mismatched_sizes=True
            ).to(device)

            if trial is not None:
                trial_metrics_cb.current_trial = trial.number
                if max_freezable > 0:
                    n_freeze = trial.suggest_int("num_frozen_layers", 0, max_freezable)
                    freeze_layers(m, model_key, n_freeze)
                else:
                    n_freeze = 0
                    freeze_layers(m, model_key, 0)
                if trial.number not in trial_records:
                    trial_records[trial.number] = {"trial_number": trial.number + 1, "params": {}}
                trial_records[trial.number]["params"]["num_frozen_layers"] = n_freeze
            else:
                freeze_layers(m, model_key, num_frozen_layers)
            return m

        def hp_space(trial):
            params = {
                "learning_rate": trial.suggest_float("learning_rate", 1e-5, 8e-5, log=True),
                "weight_decay": trial.suggest_float("weight_decay", 0.0, 0.15),
                "per_device_train_batch_size": trial.suggest_categorical("per_device_train_batch_size", [8, 16]),
                "warmup_ratio": trial.suggest_float("warmup_ratio", 0.05, 0.25),
            }
            if trial.number not in trial_records:
                trial_records[trial.number] = {"trial_number": trial.number + 1, "params": {}}
            trial_records[trial.number]["params"].update(copy.deepcopy(params))
            return params

        search_args = TrainingArguments(
            output_dir=search_dir,
            eval_strategy="epoch",
            logging_strategy="epoch",
            learning_rate=5e-5,
            per_device_train_batch_size=batch_size,
            per_device_eval_batch_size=batch_size,
            num_train_epochs=epochs,
            disable_tqdm=True,
            save_strategy="epoch",
            load_best_model_at_end=True,
            metric_for_best_model="f1",
            greater_is_better=True,
            save_total_limit=1,
            report_to="none",
        )

        search_trainer = WeightedLossTrainer(
            model_init=model_init,
            args=search_args,
            train_dataset=tokenized_train,
            eval_dataset=tokenized_test,
            processing_class=tokenizer,
            data_collator=data_collator,
            compute_metrics=compute_metrics,
            callbacks=[trial_metrics_cb],
        )
        search_trainer.remove_callback(PrinterCallback)

        best_run = search_trainer.hyperparameter_search(
            direction="maximize",
            backend="optuna",
            hp_space=hp_space,
            n_trials=n_trials,
            compute_objective=lambda m: m["eval_f1"],
        )

        best_hp.update(best_run.hyperparameters)
        best_f1 = best_run.objective

        # Build comprehensive trials dictionary from trial records and callback
        trials_dict = {}
        for trial_num, rec in trial_records.items():
            t_key = str(trial_num + 1)
            trial_bundle = trial_metrics_cb.trials_history.get(trial_num, {})
            epoch_records = trial_bundle.get("epochs", [])

            epochs_list = [e.get("Epoch") for e in epoch_records]
            f1s_list = [e.get("F1", 0.0) for e in epoch_records]
            eval_losses = [e.get("Eval_Loss", e.get("Loss", 0.0)) for e in epoch_records]
            train_losses = [e.get("Train_Loss", eval_losses[i] if i < len(eval_losses) else 0.0) for i, e in enumerate(epoch_records)]
            precs = [e.get("Precision", 0.0) for e in epoch_records]
            recs = [e.get("Recall", 0.0) for e in epoch_records]
            accs = [e.get("Accuracy", 0.0) for e in epoch_records]
            vram_allocs = [e.get("VRAM_Allocated_MB", 0.0) for e in epoch_records]
            vram_reserves = [e.get("VRAM_Reserved_MB", 0.0) for e in epoch_records]
            speeds = [e.get("Samples_Per_Sec", 0.0) for e in epoch_records if e.get("Samples_Per_Sec")]
            avg_speed = round(sum(speeds) / len(speeds), 1) if speeds else 0.0

            trial_dur = trial_bundle.get("trial_duration_sec", 0.0)
            m_dur = int(trial_dur // 60)
            s_dur = int(trial_dur % 60)
            formatted_dur = f"{m_dur}m {s_dur:02d}s" if m_dur > 0 else f"{s_dur}s"

            best_trial_f1 = max(f1s_list) if f1s_list else 0.0

            trials_dict[t_key] = {
                "trial_number": trial_num + 1,
                "state": "COMPLETE",
                "params": rec["params"],
                "best_f1": round(float(best_trial_f1), 4),
                "peak_vram_mb": trial_bundle.get("peak_vram_mb", 0.0),
                "peak_vram_gb": trial_bundle.get("peak_vram_gb", 0.0),
                "avg_samples_per_sec": avg_speed,
                "trial_duration_sec": trial_dur,
                "formatted_duration": formatted_dur,
                "epochs": epochs_list,
                "f1": f1s_list,
                "loss": eval_losses,
                "eval_loss": eval_losses,
                "train_loss": train_losses,
                "precision": precs,
                "recall": recs,
                "accuracy": accs,
                "vram_allocated_mb": vram_allocs,
                "vram_reserved_mb": vram_reserves,
            }

        if os.path.exists(search_dir):
            shutil.rmtree(search_dir)

        print(f"\nOptuna Search Finished. Best F1: {best_f1:.4f}")
        print(f"Optimal Parameters: {best_hp}")

    # Final Training Run
    print(f"\n--- Executing Final Training Run ---")
    final_frozen = best_hp.get("num_frozen_layers", num_frozen_layers)
    final_model = model_cls.from_pretrained(
        pretrained_id,
        num_labels=len(label2id),
        id2label=id2label,
        label2id=label2id,
        ignore_mismatched_sizes=True
    ).to(device)

    freeze_layers(final_model, model_key, final_frozen)

    trainable_params = sum(p.numel() for p in final_model.parameters() if p.requires_grad)
    total_params = sum(p.numel() for p in final_model.parameters())
    print(f"Model Parameters: {trainable_params:,} trainable / {total_params:,} total ({trainable_params / total_params * 100:.1f}%)")

    final_metrics_cb = ComprehensiveMetricsCallback(
        mode="final",
        total_train_samples=len(tokenized_train),
        num_epochs=epochs
    )

    final_args = TrainingArguments(
        output_dir=output_dir,
        eval_strategy="epoch",
        logging_strategy="epoch",
        learning_rate=best_hp.get("learning_rate", learning_rate),
        per_device_train_batch_size=best_hp.get("per_device_train_batch_size", batch_size),
        per_device_eval_batch_size=batch_size,
        num_train_epochs=epochs,
        weight_decay=best_hp.get("weight_decay", weight_decay),
        disable_tqdm=True,
        save_strategy="epoch",
        load_best_model_at_end=True,
        metric_for_best_model="f1",
        greater_is_better=True,
        save_total_limit=1,
        warmup_ratio=best_hp.get("warmup_ratio", warmup_ratio),
        report_to="none",
        fp16=True
    )

    final_trainer = WeightedLossTrainer(
        model=final_model,
        args=final_args,
        train_dataset=tokenized_train,
        eval_dataset=tokenized_test,
        processing_class=tokenizer,
        data_collator=data_collator,
        compute_metrics=compute_metrics,
        callbacks=[final_metrics_cb],
    )
    final_trainer.remove_callback(PrinterCallback)

    final_trainer.train()

    # Save final model and tokenizer
    print(f"Saving fine-tuned model and tokenizer to {output_dir}...")
    final_trainer.save_model(output_dir)
    tokenizer.save_pretrained(output_dir)

    # Hold-out test set evaluation
    print("\n--- Evaluating Best Model Checkpoint on 100% Human Hold-Out Test Split ---")
    test_results = final_trainer.evaluate(eval_dataset=tokenized_test)
    test_eval = {
        "f1": round(float(test_results.get("eval_f1", 0.0)), 4),
        "precision": round(float(test_results.get("eval_precision", 0.0)), 4),
        "recall": round(float(test_results.get("eval_recall", 0.0)), 4),
        "loss": round(float(test_results.get("eval_loss", 0.0)), 4),
        "accuracy": round(float(test_results.get("eval_accuracy", 0.0)), 4),
    }
    print(f"Hold-Out Test Score: F1={test_eval['f1']} | Recall={test_eval['recall']} | Precision={test_eval['precision']} | Loss={test_eval['loss']}")

    # Assemble comprehensive metrics
    history = final_metrics_cb.get_final_history_list()
    final_epochs = [h["Epoch"] for h in history]
    final_f1 = [h.get("F1", 0.0) for h in history]
    final_eval_loss = [h.get("Eval_Loss", 0.0) for h in history]
    final_train_loss = [h.get("Train_Loss", final_eval_loss[idx] if idx < len(final_eval_loss) else 0.0) for idx, h in enumerate(history)]
    final_prec = [h.get("Precision", 0.0) for h in history]
    final_rec = [h.get("Recall", 0.0) for h in history]
    final_acc = [h.get("Accuracy", 0.0) for h in history]
    final_lr = [h.get("Learning_Rate", 0.0) for h in history]

    peak_f1 = max(final_f1) if final_f1 else best_f1
    peak_epoch = final_epochs[final_f1.index(peak_f1)] if final_f1 else None

    final_dur = round(time.time() - (final_metrics_cb.train_start_time or time.time()), 1)
    fm_dur = int(final_dur // 60)
    fs_dur = int(final_dur % 60)
    final_dur_str = f"{fm_dur}m {fs_dur:02d}s" if fm_dur > 0 else f"{fs_dur}s"
    final_speeds = [h.get("Samples_Per_Sec", 0.0) for h in history if h.get("Samples_Per_Sec")]
    avg_final_speed = round(sum(final_speeds) / len(final_speeds), 1) if final_speeds else 0.0

    print("\n" + "=" * 65)
    print(f"  TRAINING TELEMETRY SUMMARY: {model_key.upper()}")
    print("=" * 65)
    if torch.cuda.is_available():
        gpu_name = torch.cuda.get_device_name(0)
        tot_vram = torch.cuda.get_device_properties(0).total_memory / (1024 ** 3)
        peak_alloc = final_metrics_cb.peak_vram_allocated_mb
        peak_res = final_metrics_cb.peak_vram_reserved_mb
        util = (peak_res / (tot_vram * 1024)) * 100
        print(f"  Compute Device:          {gpu_name} ({tot_vram:.2f} GB VRAM)")
        print(f"  Peak VRAM Allocated:     {peak_alloc:,.1f} MB ({peak_alloc / 1024:.2f} GB)")
        print(f"  Peak VRAM Reserved:      {peak_res:,.1f} MB ({peak_res / 1024:.2f} GB, {util:.1f}% capacity)")
    else:
        print(f"  Compute Device:          CPU")
    if psutil is not None:
        print(f"  Peak Process RAM:        {final_metrics_cb.peak_process_ram_mb:,.1f} MB")
    print(f"  Total Training Duration: {final_dur_str}")
    print(f"  Average Epoch Duration:  {final_dur / max(1, len(final_epochs)):.1f}s")
    print(f"  Average Throughput:      {avg_final_speed} samples/sec")
    print("=" * 65)

    metrics_data = {
        "model_info": {
            "name": model_key,
            "pretrained_id": pretrained_id,
            "total_parameters": total_params,
            "trainable_parameters": trainable_params,
            "device": device,
            "device_name": torch.cuda.get_device_name(0) if torch.cuda.is_available() else "CPU",
            "timestamp": datetime.datetime.now().isoformat()
        },
        "best_hp": {
            "num_frozen_layers": final_frozen,
            "learning_rate": best_hp.get("learning_rate"),
            "weight_decay": best_hp.get("weight_decay"),
            "batch_size": best_hp.get("per_device_train_batch_size", batch_size),
            "warmup_ratio": best_hp.get("warmup_ratio"),
            "best_f1": round(float(peak_f1), 4)
        },
        "hardware_telemetry": {
            "device_name": torch.cuda.get_device_name(0) if torch.cuda.is_available() else "CPU",
            "cuda_available": torch.cuda.is_available(),
            "total_vram_gb": round(torch.cuda.get_device_properties(0).total_memory / (1024 ** 3), 2) if torch.cuda.is_available() else 0.0,
            "peak_vram_allocated_mb": round(final_metrics_cb.peak_vram_allocated_mb, 1),
            "peak_vram_allocated_gb": round(final_metrics_cb.peak_vram_allocated_mb / 1024, 2),
            "peak_vram_reserved_mb": round(final_metrics_cb.peak_vram_reserved_mb, 1),
            "peak_vram_reserved_gb": round(final_metrics_cb.peak_vram_reserved_mb / 1024, 2),
            "vram_utilization_pct": round((final_metrics_cb.peak_vram_reserved_mb / (torch.cuda.get_device_properties(0).total_memory / (1024 ** 2))) * 100, 1) if torch.cuda.is_available() else 0.0,
            "total_system_ram_gb": round(psutil.virtual_memory().total / (1024 ** 3), 2) if psutil else 0.0,
            "peak_process_ram_mb": round(final_metrics_cb.peak_process_ram_mb, 1)
        },
        "performance_telemetry": {
            "total_train_samples": len(tokenized_train),
            "total_test_samples": len(tokenized_test),
            "total_epochs": epochs,
            "total_duration_sec": final_dur,
            "formatted_duration": final_dur_str,
            "avg_samples_per_sec": avg_final_speed,
            "epoch_durations_sec": [h.get("Epoch_Duration_Sec", 0.0) for h in history]
        },
        "trials": trials_dict,
        "final_run": {
            "epochs": final_epochs,
            "f1": final_f1,
            "loss": final_eval_loss,
            "eval_loss": final_eval_loss,
            "train_loss": final_train_loss,
            "precision": final_prec,
            "recall": final_rec,
            "accuracy": final_acc,
            "learning_rate": final_lr,
            "vram_allocated_mb": [h.get("VRAM_Allocated_MB", 0.0) for h in history],
            "vram_reserved_mb": [h.get("VRAM_Reserved_MB", 0.0) for h in history],
            "vram_peak_gb": [h.get("VRAM_Peak_GB", 0.0) for h in history],
            "epoch_duration_sec": [h.get("Epoch_Duration_Sec", 0.0) for h in history],
            "samples_per_sec": [h.get("Samples_Per_Sec", 0.0) for h in history]
        },
        "test_eval": test_eval,
        "summary": {
            "peak_val_f1": round(float(peak_f1), 4),
            "peak_val_epoch": peak_epoch,
            "optuna_best_f1": round(float(best_f1), 4) if best_f1 else round(float(peak_f1), 4),
            "test_f1": test_eval["f1"],
            "test_recall": test_eval["recall"],
            "test_precision": test_eval["precision"],
            "total_trials": len(trials_dict)
        }
    }

    with open(metrics_file, "w", encoding="utf-8") as f:
        json.dump(metrics_data, f, indent=2)
    print(f"Comprehensive metrics ledger saved to: {metrics_file}")

    # Remove intermediate checkpoint directory to keep disk space lean
    for item in os.listdir(output_dir):
        subpath = os.path.join(output_dir, item)
        if os.path.isdir(subpath) and item.startswith("checkpoint-"):
            shutil.rmtree(subpath)

    print(f"\n[SUCCESS] {model_key.upper()} trained successfully! Peak Val F1: {peak_f1:.4f} | Test F1: {test_eval['f1']:.4f}")
    return peak_f1


def main():
    parser = argparse.ArgumentParser(description="Gotcha Clause Extractor: Production Training & Evaluation Suite")
    parser.add_argument("--model", type=str, default="electra-small", choices=AVAILABLE_MODELS + ["all"],
                        help="Model key to train, or 'all' to train all 4 models sequentially.")
    parser.add_argument("--epochs", type=int, default=10, help="Number of training epochs (default: 10).")
    parser.add_argument("--batch_size", type=int, default=8, help="Batch size (default: 8).")
    parser.add_argument("--lr", type=float, default=5e-5, help="Learning rate (default: 5e-5).")
    parser.add_argument("--optuna", action="store_true", help="Enable Optuna hyperparameter optimization.")
    parser.add_argument("--n_trials", type=int, default=10, help="Number of Optuna trials (default: 10).")
    parser.add_argument("--dry_run", action="store_true", help="Quick 1-minute dry run on tiny sample to verify pipeline.")
    parser.add_argument("--evaluate_all", action="store_true", help="Evaluate all existing models on test set and write metrics.")
    parser.add_argument("--output_dir", type=str, default="./gotcha-extractor-model", help="Base output directory.")
    parser.add_argument("--seed", type=int, default=42, help="Random seed.")
    parser.add_argument("--synthetic_path", type=str, default="corpus/synthetic_gotchas.json", help="Path to synthetic JSON.")
    args = parser.parse_args()

    device = "cuda" if torch.cuda.is_available() else "cpu"
    print("\n" + "=" * 65)
    print("  HARDWARE & SYSTEM TELEMETRY INITIALIZATION")
    print("=" * 65)
    if torch.cuda.is_available():
        gpu_name = torch.cuda.get_device_name(0)
        props = torch.cuda.get_device_properties(0)
        total_vram_gb = props.total_memory / (1024 ** 3)
        cuda_arch = f"{props.major}.{props.minor}"
        print(f"  Compute Device:     {gpu_name} (Arch: sm_{cuda_arch})")
        print(f"  Total GPU VRAM:     {total_vram_gb:.2f} GB ({props.total_memory // (1024 ** 2):,} MB)")
        print(f"  PyTorch CUDA:       {torch.version.cuda}")
    else:
        print(f"  Compute Device:     CPU (CUDA not available)")

    if psutil is not None:
        vm = psutil.virtual_memory()
        print(f"  System Total RAM:   {vm.total / (1024 ** 3):.2f} GB ({vm.percent}% in use)")
        print(f"  CPU Cores:          {psutil.cpu_count(logical=True)} logical ({psutil.cpu_count(logical=False)} physical)")
    print("=" * 65)

    # Build dataset
    splits = build_unified_dataset(synthetic_path=args.synthetic_path, seed=args.seed)

    # If evaluate_all requested, evaluate all existing models and exit
    if args.evaluate_all:
        print("\n" + "=" * 65)
        print("  RUNNING COMPREHENSIVE HOLD-OUT TEST EVALUATION ON ALL MODELS")
        print("=" * 65)
        for m in AVAILABLE_MODELS:
            evaluate_model_on_test(m, splits, output_base_dir=args.output_dir, device=device)
        print("\nAll model test evaluations complete.")
        return

    if args.dry_run:
        print("\n*** DRY RUN MODE: Truncating splits to 16 train / 8 test samples for fast verification ***")
        args.output_dir = "./tmp_dry_run"
        splits["train"] = splits["train"].select(range(min(16, len(splits["train"]))))
        splits["test"] = splits["test"].select(range(min(8, len(splits["test"]))))
        args.epochs = 2
        if args.optuna:
            args.n_trials = 2

    models_to_train = AVAILABLE_MODELS if args.model == "all" else [args.model]

    for m in models_to_train:
        train_single_model(
            model_key=m,
            dataset_splits=splits,
            output_base_dir=args.output_dir,
            epochs=args.epochs,
            batch_size=args.batch_size,
            learning_rate=args.lr,
            use_optuna=args.optuna,
            n_trials=args.n_trials,
            device=device
        )

    print("\n" + "=" * 65)
    print("  ALL REQUESTED TRAINING & EVALUATION COMPLETED SUCCESSFULLY!")
    print("=" * 65)


if __name__ == "__main__":
    main()