model/train.py

"""
IndoBERT Training Script with Stratified 5-Fold Cross-Validation.
Fine-tunes IndoBERT for essay scoring (1-5 scale).

Pair Encoding Mode:
- Set PAIR_ENCODING=True di config.py untuk menggunakan pair encoding
  ([CLS] jawaban_siswa [SEP] kunci_jawaban [SEP])
- Mode pair: load dataset_indonesia_pair.csv + kunci_jawaban column
- Model disimpan ke best_model_pair/ (tidak overwrite model lama)
- Mode single-text (lama) tetap berfungsi dengan PAIR_ENCODING=False

Usage:
    python -m model.train
"""
import os
import sys
sys.stdout.reconfigure(encoding='utf-8', errors='replace')
import json
import numpy as np
import pandas as pd
import torch
from torch.utils.data import DataLoader
from transformers import (
    AutoTokenizer,
    AutoModelForSequenceClassification,
    get_linear_schedule_with_warmup,
)
from sklearn.model_selection import StratifiedKFold
from sklearn.utils.class_weight import compute_class_weight
from tqdm import tqdm

from model.config import (
    DATASET_PATH, TEXT_COLUMN, LABEL_COLUMN, ID_SOAL_COLUMN, NUM_LABELS,
    INDOBERT_MODEL_NAME, MAX_SEQ_LENGTH, BATCH_SIZE,
    GRADIENT_ACCUMULATION_STEPS, LEARNING_RATE, NUM_EPOCHS,
    WARMUP_RATIO, WEIGHT_DECAY, EARLY_STOPPING_PATIENCE,
    NUM_FOLDS, SEED, FP16, SAVED_MODELS_DIR,
    # BARU: Pair encoding config
    PAIR_ENCODING, DATASET_PATH_PAIR, SAVED_MODELS_DIR_PAIR, KUNCI_JAWABAN_COLUMN,
    FOCAL_LOSS, FOCAL_LOSS_GAMMA,
    SAVED_MODELS_DIR_FOCAL, SAVED_MODELS_DIR_PAIR_FOCAL,
)
from model.dataset import EssayDataset
from model.evaluate import compute_all_metrics, print_evaluation_report
from model.loss import FocalLoss


def set_seed(seed):
    """Set random seeds for reproducibility."""
    torch.manual_seed(seed)
    np.random.seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)


def train_one_epoch(model, loader, optimizer, scheduler, criterion, device, scaler, grad_accum_steps):
    """Train for one epoch."""
    model.train()
    total_loss = 0.0
    optimizer.zero_grad()

    for step, batch in enumerate(tqdm(loader, desc="  Training", leave=False)):
        input_ids = batch["input_ids"].to(device)
        attention_mask = batch["attention_mask"].to(device)
        labels = batch["labels"].to(device)

        if scaler is not None:
            with torch.amp.autocast("cuda"):
                outputs = model(input_ids=input_ids, attention_mask=attention_mask)
                loss = criterion(outputs.logits, labels)
                loss = loss / grad_accum_steps
            scaler.scale(loss).backward()

            if (step + 1) % grad_accum_steps == 0:
                scaler.step(optimizer)
                scaler.update()
                scheduler.step()
                optimizer.zero_grad()
        else:
            outputs = model(input_ids=input_ids, attention_mask=attention_mask)
            loss = criterion(outputs.logits, labels)
            loss = loss / grad_accum_steps
            loss.backward()

            if (step + 1) % grad_accum_steps == 0:
                optimizer.step()
                scheduler.step()
                optimizer.zero_grad()

        total_loss += loss.item() * grad_accum_steps

    return total_loss / len(loader)


def evaluate_model(model, loader, device):
    """Evaluate model and return predictions + labels."""
    model.eval()
    all_preds = []
    all_labels = []

    with torch.no_grad():
        for batch in tqdm(loader, desc="  Evaluating", leave=False):
            input_ids = batch["input_ids"].to(device)
            attention_mask = batch["attention_mask"].to(device)
            labels = batch["labels"].to(device)

            outputs = model(input_ids=input_ids, attention_mask=attention_mask)
            preds = torch.argmax(outputs.logits, dim=-1)

            all_preds.extend(preds.cpu().numpy())
            all_labels.extend(labels.cpu().numpy())

    # Convert 0-4 back to 1-5
    all_preds = [p + 1 for p in all_preds]
    all_labels = [l + 1 for l in all_labels]

    return all_labels, all_preds


def train():
    """Main training loop with 5-fold cross-validation."""
    print("=" * 60)
    print("  AES-Feedback: IndoBERT Training")
    print("=" * 60)

    set_seed(SEED)

    # ── Device ──────────────────────────────────────────
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"\n  Device: {device}")
    if torch.cuda.is_available():
        print(f"  GPU: {torch.cuda.get_device_name(0)}")
        print(f"  VRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB")

    # ── Load Data ───────────────────────────────────────
    # BARU: Pilih dataset berdasarkan mode pair encoding
    if PAIR_ENCODING:
        dataset_path = DATASET_PATH_PAIR
        print(f"\n  [PAIR ENCODING] Loading dataset: {dataset_path}")
    else:
        dataset_path = DATASET_PATH
        print(f"\n  Loading dataset: {dataset_path}")

    df = pd.read_csv(dataset_path)
    texts = df[TEXT_COLUMN].values
    labels = df[LABEL_COLUMN].values
    id_soal = df[ID_SOAL_COLUMN].values

    # BARU: Load kunci_jawaban jika pair encoding
    kunci_jawaban = None
    if PAIR_ENCODING:
        kunci_jawaban = df[KUNCI_JAWABAN_COLUMN].values
        print(f"  Kunci jawaban loaded: {len(kunci_jawaban)} entries")

    print(f"  Samples: {len(df)}")
    print(f"  Score distribution: {dict(pd.Series(labels).value_counts().sort_index())}")
    print(f"  id_soal distribution: {dict(pd.Series(id_soal).value_counts())}")

    # Combined stratify key: ensure proportional (id_soal, score) per fold
    stratify_key = pd.Series(id_soal).astype(str) + "_" + pd.Series(labels).astype(str)

    # ── Model save directory ──────────────────────────
    # Pilih direktori berdasarkan (PAIR_ENCODING, FOCAL_LOSS)
    if FOCAL_LOSS and PAIR_ENCODING:
        models_root = SAVED_MODELS_DIR_PAIR_FOCAL
        mode_label = "PAIR ENCODING + FOCAL LOSS"
    elif FOCAL_LOSS:
        models_root = SAVED_MODELS_DIR_FOCAL
        mode_label = "FOCAL LOSS"
    elif PAIR_ENCODING:
        models_root = SAVED_MODELS_DIR_PAIR
        mode_label = "PAIR ENCODING"
    else:
        models_root = SAVED_MODELS_DIR
        mode_label = "SINGLE TEXT"
    print(f"\n  [{mode_label}] Model akan disimpan ke: {models_root}")

    # ── Tokenizer ───────────────────────────────────────
    print(f"\n  Loading tokenizer: {INDOBERT_MODEL_NAME}")
    tokenizer = AutoTokenizer.from_pretrained(INDOBERT_MODEL_NAME)

    # ── Cross-Validation ────────────────────────────────
    skf = StratifiedKFold(n_splits=NUM_FOLDS, shuffle=True, random_state=SEED)
    all_fold_metrics = []
    best_overall_qwk = -1.0
    loss_label = "Focal Loss" if FOCAL_LOSS else "CrossEntropy"

    for fold_idx, (train_idx, val_idx) in enumerate(skf.split(texts, stratify_key)):
        print(f"\n{'='*60}")
        print(f"  FOLD {fold_idx + 1}/{NUM_FOLDS}  ({loss_label})")
        print(f"  Train: {len(train_idx)} | Val: {len(val_idx)}")
        train_soal_dist = dict(pd.Series(id_soal[train_idx]).value_counts())
        val_soal_dist = dict(pd.Series(id_soal[val_idx]).value_counts())
        print(f"  Train per soal: {train_soal_dist}")
        print(f"  Val per soal:   {val_soal_dist}")
        print(f"{'='*60}")

        # Create datasets
        # BARU: Jika PAIR_ENCODING, kirim kunci_jawaban sebagai text_pair
        if PAIR_ENCODING:
            train_dataset = EssayDataset(
                texts[train_idx], labels[train_idx], tokenizer, MAX_SEQ_LENGTH,
                text_pair=kunci_jawaban[train_idx],
            )
            val_dataset = EssayDataset(
                texts[val_idx], labels[val_idx], tokenizer, MAX_SEQ_LENGTH,
                text_pair=kunci_jawaban[val_idx],
            )
        else:
            train_dataset = EssayDataset(
                texts[train_idx], labels[train_idx], tokenizer, MAX_SEQ_LENGTH
            )
            val_dataset = EssayDataset(
                texts[val_idx], labels[val_idx], tokenizer, MAX_SEQ_LENGTH
            )

        train_loader = DataLoader(
            train_dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=0, pin_memory=True
        )
        val_loader = DataLoader(
            val_dataset, batch_size=BATCH_SIZE, shuffle=False, num_workers=0, pin_memory=True
        )

        # ── Model ───────────────────────────────────────
        model = AutoModelForSequenceClassification.from_pretrained(
            INDOBERT_MODEL_NAME, num_labels=NUM_LABELS
        )
        model.to(device)

        # ── Class Weights & Loss ────────────────────────
        class_weights = compute_class_weight(
            "balanced", classes=np.arange(NUM_LABELS), y=labels[train_idx] - 1
        )
        class_weights_tensor = torch.tensor(class_weights, dtype=torch.float).to(device)

        if FOCAL_LOSS:
            criterion = FocalLoss(gamma=FOCAL_LOSS_GAMMA, weight=class_weights_tensor)
            loss_label = f"Focal Loss (gamma={FOCAL_LOSS_GAMMA})"
        else:
            criterion = torch.nn.CrossEntropyLoss(weight=class_weights_tensor)
            loss_label = "CrossEntropy (balanced)"

        # ── Optimizer & Scheduler ───────────────────────
        optimizer = torch.optim.AdamW(
            model.parameters(), lr=LEARNING_RATE, weight_decay=WEIGHT_DECAY
        )

        total_steps = (len(train_loader) * NUM_EPOCHS) // GRADIENT_ACCUMULATION_STEPS
        warmup_steps = int(total_steps * WARMUP_RATIO)
        scheduler = get_linear_schedule_with_warmup(
            optimizer, num_warmup_steps=warmup_steps, num_training_steps=total_steps
        )

        # ── Mixed Precision ─────────────────────────────
        scaler = None
        if FP16 and torch.cuda.is_available():
            scaler = torch.amp.GradScaler("cuda")

        # ── Training Loop ───────────────────────────────
        best_fold_qwk = -1.0
        patience_counter = 0

        for epoch in range(NUM_EPOCHS):
            # Train
            avg_loss = train_one_epoch(
                model, train_loader, optimizer, scheduler, criterion,
                device, scaler, GRADIENT_ACCUMULATION_STEPS
            )

            # Evaluate
            y_true, y_pred = evaluate_model(model, val_loader, device)
            metrics = compute_all_metrics(y_true, y_pred)

            print(
                f"  Epoch {epoch+1:>2}/{NUM_EPOCHS} │ "
                f"Loss: {avg_loss:.4f} │ "
                f"QWK: {metrics['qwk']:.4f} │ "
                f"Acc: {metrics['accuracy']:.4f} │ "
                f"F1: {metrics['f1_macro']:.4f}"
            )

            # Save best model for this fold
            if metrics["qwk"] > best_fold_qwk:
                best_fold_qwk = metrics["qwk"]
                patience_counter = 0

                # Save fold model
                fold_dir = os.path.join(models_root, f"fold_{fold_idx}")
                os.makedirs(fold_dir, exist_ok=True)
                model.save_pretrained(fold_dir)
                tokenizer.save_pretrained(fold_dir)

                # Save overall best
                if metrics["qwk"] > best_overall_qwk:
                    best_overall_qwk = metrics["qwk"]
                    best_dir = os.path.join(models_root, "best_model")
                    os.makedirs(best_dir, exist_ok=True)
                    model.save_pretrained(best_dir)
                    tokenizer.save_pretrained(best_dir)
                    print(f"    ★ New best model saved! (QWK: {metrics['qwk']:.4f})")
            else:
                patience_counter += 1
                if patience_counter >= EARLY_STOPPING_PATIENCE:
                    print(f"    ⏹ Early stopping at epoch {epoch+1}")
                    break

        # ── Final Fold Evaluation ───────────────────────
        # Reload best fold model
        best_fold_model = AutoModelForSequenceClassification.from_pretrained(
            os.path.join(models_root, f"fold_{fold_idx}"),
            num_labels=NUM_LABELS,
        )
        best_fold_model.to(device)

        y_true, y_pred = evaluate_model(best_fold_model, val_loader, device)
        fold_metrics = compute_all_metrics(y_true, y_pred)
        all_fold_metrics.append(fold_metrics)

        print(f"\n  Fold {fold_idx+1} Best Results:")
        for name, value in fold_metrics.items():
            print(f"    {name:>15}: {value:.4f}")

        # Clean up GPU memory
        del model, best_fold_model
        torch.cuda.empty_cache() if torch.cuda.is_available() else None

    # ── Final Cross-Validation Summary ──────────────────
    print(f"\n{'='*60}")
    print("  FINAL 5-FOLD CROSS-VALIDATION RESULTS")
    print(f"{'='*60}")

    final_results = {}
    for metric_name in all_fold_metrics[0].keys():
        values = [m[metric_name] for m in all_fold_metrics]
        mean_val = float(np.mean(values))
        std_val = float(np.std(values))
        final_results[metric_name] = {"mean": mean_val, "std": std_val}
        print(f"  {metric_name:>15}: {mean_val:.4f} ± {std_val:.4f}")

    print(f"\n  Best Overall QWK: {best_overall_qwk:.4f}")
    print(f"  Best model saved to: {os.path.join(models_root, 'best_model')}")

    # Save results to JSON
    results_suffix = ""
    if PAIR_ENCODING:
        results_suffix += "_pair"
    if FOCAL_LOSS:
        results_suffix += "_focal"
    results_path = os.path.join(models_root, f"cv_results{results_suffix}.json")
    os.makedirs(models_root, exist_ok=True)
    with open(results_path, "w") as f:
        json.dump(
            {
                "fold_metrics": [
                    {k: float(v) for k, v in m.items()} for m in all_fold_metrics
                ],
                "summary": {
                    k: {"mean": v["mean"], "std": v["std"]}
                    for k, v in final_results.items()
                },
                "best_qwk": float(best_overall_qwk),
                "config": {
                    "model": INDOBERT_MODEL_NAME,
                    "num_folds": NUM_FOLDS,
                    "max_seq_length": MAX_SEQ_LENGTH,
                    "batch_size": BATCH_SIZE,
                    "learning_rate": LEARNING_RATE,
                    "num_epochs": NUM_EPOCHS,
                    "seed": SEED,
                    "pair_encoding": PAIR_ENCODING,
                    "focal_loss": FOCAL_LOSS,
                    "focal_loss_gamma": FOCAL_LOSS_GAMMA if FOCAL_LOSS else None,
                },
            },
            f,
            indent=2,
        )
    print(f"  Results saved to: {results_path}")
    print("\n  ✓ Training complete!")


if __name__ == "__main__":
    train()