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