RewardModel.py

import math
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import random
from tqdm import tqdm
from torch.utils.data import Dataset, DataLoader
from transformers import AutoTokenizer, AutoModel, AdamW, get_linear_schedule_with_warmup
from torch.cuda.amp import GradScaler
from contextlib import nullcontext
from typing import List, Dict, Optional
from pairadigm import Pairadigm
import json

# Set random seeds for reproducibility
def set_seeds(seed: int = 42):
    np.random.seed(seed)
    torch.manual_seed(seed)
    random.seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False

class RewardModel(nn.Module):
    """
    Unified class for training and using a reward model for text scoring.
    
    This class handles:
    - Model initialization and configuration
    - Dataset creation and management
    - Training loop with pairwise comparisons
    - Scoring individual texts or batches
    - Score normalization
    """
    
    def __init__(
        self,
        model_name: str = "roberta-large",
        dropout: float = 0.1,
        max_length: int = 384,
        device: Optional[str] = None,
        Pairadigm: Optional['Pairadigm'] = None,
        seed: int = 42
    ):
        """
        Initialize the reward model trainer.
        
        Args:
            model_name: HuggingFace model identifier
            dropout: Dropout rate for reward head
            max_length: Maximum sequence length for tokenization
            device: Device to use ('cuda', 'cpu', or None for auto-detect)
        """
        super().__init__()
        # Set seeds at model initialization
        set_seeds(seed)

        self.model_name = model_name
        self.max_length = max_length
        if device:
            self.device = device
        elif torch.cuda.is_available():
            self.device = 'cuda'
        elif torch.backends.mps.is_available():
            self.device = 'mps'
        else:
            self.device = 'cpu'
        print(f"Model using device: {self.device}")
        self.device_type = self.device.type if isinstance(self.device, torch.device) else self.device.split(':')[0]
        self.use_amp = self.device_type == 'cuda'

        # Initialize tokenizer
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        
        # Initialize model
        self.model = self._build_model(dropout)
        self.model.to(self.device)

        # If a Pairadigm instance is provided, link it
        self.pairadigm = Pairadigm
        
        # Training state
        self.optimizer = None
        self.scheduler = None
        #self.scaler = GradScaler('cuda') # For mixed precision training
        self.scaler = GradScaler(self.device_type, enabled=self.device_type == 'cuda')
        
        self.training_history = []
    
    def _autocast_context(self):
        if self.device_type == "cuda":
            return torch.amp.autocast(device_type="cuda", dtype=torch.float16)
        if self.device_type == "cpu":
            return torch.amp.autocast(device_type="cpu", dtype=torch.bfloat16)
        return nullcontext()

    def _build_model(self, dropout: float):
        class _EncoderWithHead(nn.Module):
            
            def __init__(inner_self, model_name, dropout):
                super().__init__()
                
                inner_self.encoder = AutoModel.from_pretrained(model_name)
                hidden_size = inner_self.encoder.config.hidden_size
                inner_self.dropout = nn.Dropout(dropout)

                inner_self.valence_head = nn.Linear(hidden_size, 1)
                inner_self.arousal_dominance_head = nn.Linear(hidden_size, 2)

                nn.init.xavier_uniform_(inner_self.valence_head.weight, gain=1.0)
                nn.init.zeros_(inner_self.valence_head.bias)
                nn.init.xavier_uniform_(inner_self.arousal_dominance_head.weight, gain=1.0)
                nn.init.zeros_(inner_self.arousal_dominance_head.bias)

            def forward(inner_self, input_ids, attention_mask):
                outputs = inner_self.encoder(input_ids=input_ids, attention_mask=attention_mask)
                
                pooled_output = outputs.last_hidden_state[:, 0, :]
                pooled_output = inner_self.dropout(pooled_output)

                valence = inner_self.valence_head(pooled_output)
                arousal_dominance = inner_self.arousal_dominance_head(pooled_output)

                return valence, arousal_dominance

        return _EncoderWithHead(self.model_name, dropout)
    
    class _PairwiseDataset(Dataset):
        """Internal dataset class for pairs with V gold scores and A/D probabilities."""
        
        def __init__(self, pairs, tokenizer, max_length):
            self.pairs = pairs
            self.tokenizer = tokenizer
            self.max_length = max_length
        
        def __len__(self):
            return len(self.pairs)
        
        def __getitem__(self, idx):
            text_A, text_B, V_A, V_B, label_A, label_D = self.pairs[idx]
            
            encoding_A = self.tokenizer(
                text_A,
                max_length=self.max_length,
                padding='max_length',
                truncation=True,
                return_tensors='pt'
            )
            
            encoding_B = self.tokenizer(
                text_B,
                max_length=self.max_length,
                padding='max_length',
                truncation=True,
                return_tensors='pt'
            )
            
            return {
                'input_ids_A': encoding_A['input_ids'].squeeze(0),
                'attention_mask_A': encoding_A['attention_mask'].squeeze(0),
                'input_ids_B': encoding_B['input_ids'].squeeze(0),
                'attention_mask_B': encoding_B['attention_mask'].squeeze(0),
                'V_A': torch.tensor(V_A, dtype=torch.float),
                'V_B': torch.tensor(V_B, dtype=torch.float),
                'label_A': torch.tensor(label_A, dtype=torch.float), # 1.0 if item1 wins, 0.0 if item2 wins
                'label_D': torch.tensor(label_D, dtype=torch.float)
            }
    
    def fit(
        self,
        train_loader: DataLoader,
        eval_loader: DataLoader,
        epochs: int = 5,
        learning_rate: float = 2e-5,
        weight_decay: float = 0.01,
        warmup_steps: int = 100,
        log_interval: int = 50,
        early_stopping_patience: int = 3,
        accumulation_steps: int = 1,
        max_steps_per_epoch: int = None
    ):
        """
        Train the reward model on pairwise comparison data with optional early stopping and gradient accumulation.

        Args:
            train_loader: DataLoader with training pairs (mini-batch size)
            eval_loader: DataLoader for evaluation (required for early stopping)
            epochs: Number of training epochs
            learning_rate: Learning rate for optimizer
            warmup_steps: Number of warmup steps for scheduler
            log_interval: Log metrics every N steps
            early_stopping_patience: Number of epochs with no improvement on eval loss before stopping early.
                                     Set to None or 0 to disable early stopping.
            accumulation_steps: Number of gradient accumulation steps (effective_batch_size = mini_batch_size * accumulation_steps)
        Returns:
            The model (self.model) restored to the best-performing weights observed on eval data.
        """
        self.model.train()
        
        # Setup optimizer and scheduler
        self.optimizer = AdamW(
            self.model.parameters(), 
            lr=learning_rate, 
            weight_decay=weight_decay
        )

        steps_per_epoch = math.ceil(len(train_loader) / accumulation_steps)
        if max_steps_per_epoch is not None:
            steps_per_epoch = math.ceil(max_steps_per_epoch / accumulation_steps)

        total_optimization_steps = steps_per_epoch * epochs

        self.scheduler = get_linear_schedule_with_warmup(
            self.optimizer,
            num_warmup_steps=warmup_steps,
            num_training_steps=total_optimization_steps
        )

        # Early stopping bookkeeping
        best_state = None
        best_eval_loss = float('inf')
        epochs_without_improve = 0
        use_early_stopping = bool(early_stopping_patience and eval_loader is not None and early_stopping_patience > 0)
        
        for epoch in range(epochs):
        
            epoch_loss = 0
            grad_norms = []

            progress_bar = tqdm(train_loader, desc=f"Epoch {epoch + 1}/{epochs}")
            grad_norm = 0.0
            
            for step, batch in enumerate(progress_bar):
                # Perform forward/backward pass, accumulate gradients
                loss = self._training_step(batch, accumulation_steps)
                epoch_loss += loss

                # Check our conditions for stepping the optimizer
                is_accum_step = (step + 1) % accumulation_steps == 0
                is_last_batch = (step + 1) == len(train_loader)
                is_max_step = max_steps_per_epoch is not None and (step + 1) == max_steps_per_epoch

                if is_accum_step or is_last_batch or is_max_step:
                    self.scaler.unscale_(self.optimizer)
                    grad_norm = torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0)
                    grad_norms.append(grad_norm.item())
                    
                    scale_before = self.scaler.get_scale()
                    self.scaler.step(self.optimizer)
                    self.scaler.update()
                    if self.scaler.get_scale() >= scale_before:  # no NaN/Inf — step actually happened
                        self.scheduler.step()
                    self.optimizer.zero_grad(set_to_none=True)
                
                if (step + 1) % log_interval == 0:
                    avg_loss = epoch_loss / (step + 1)
                    progress_bar.set_postfix({'loss': f'{avg_loss:.4f}', 'grad_norm': f'{grad_norm:.4f}'})
                
                # Break AFTER processing the step, so everything stays perfectly aligned
                if is_max_step:
                    break

            # The loop is done. No extra cleanup block needed!
            actual_steps = step + 1
            avg_epoch_loss = epoch_loss / actual_steps
            epoch_metrics = {
                'epoch': epoch + 1, 
                'train_loss': avg_epoch_loss,
                'train_grad_norm': float(np.nanmean([g for g in grad_norms if np.isfinite(g)])) if grad_norms else float('nan'),
                'grad_norms': grad_norms 
            }
            
            print(f"Epoch {epoch + 1} - Train Loss: {avg_epoch_loss:.4f}")
            
            # Evaluation 
            if eval_loader:
                eval_metrics = self.evaluate(eval_loader)
                epoch_metrics.update(eval_metrics)

                print(f"Epoch {epoch + 1} - Eval Loss: {eval_metrics['eval_loss']:.4f} | Accuracy: {eval_metrics['eval_accuracy']:.2%}")

                print(f"  Loss - V (MSE): {eval_metrics['eval_loss_V']:.4f} | A: {eval_metrics['eval_loss_A']:.4f} | D: {eval_metrics['eval_loss_D']:.4f}")

                print(f"  Acc  - V (EMOBANK): {eval_metrics['eval_accuracy_V']:.2%} | A: {eval_metrics['eval_accuracy_A']:.2%} | D: {eval_metrics['eval_accuracy_D']:.2%}")
                
                # Check for improvement on eval_loss and save best model
                current_eval_loss = eval_metrics.get('eval_loss', float('inf'))
                if current_eval_loss < best_eval_loss:
                    best_eval_loss = current_eval_loss
                    # store a CPU copy of the state dict
                    best_state = {k: v.cpu().clone() for k, v in self.model.state_dict().items()}
                    epochs_without_improve = 0
                    print(f"  New best model found (eval_loss improved to {best_eval_loss:.4f}).")
                else:
                    epochs_without_improve += 1
                    print(f"  No improvement for {epochs_without_improve} epoch(s).")
            
            else:
                # If no eval_loader provided, we can't do early stopping / track best by eval
                best_state = {k: v.cpu().clone() for k, v in self.model.state_dict().items()}
            
            self.training_history.append(epoch_metrics)

            # Early stopping check
            if use_early_stopping and epochs_without_improve >= early_stopping_patience:
                print(f"Early stopping triggered after {epoch + 1} epochs (no improvement in eval loss for {early_stopping_patience} epochs).")
                break
        
        # Restore best model weights if we tracked them
        if best_state is not None:
            # move tensors back to device as needed when loading
            device_state = {k: v.to(self.device) for k, v in best_state.items()}
            self.model.load_state_dict(device_state)
            print("Best model weights restored based on eval data.")
        
        return self.model
    
    def _training_step(self, batch, accumulation_steps: int = 1) -> float:
        """Single training step: MSE for V (pointwise) and binary cross-entropy for A/D (pairwise).
        
        Args:
            batch: Dictionary containing input_ids, attention_masks,
                   V_A/V_B (gold scores), label_A/label_D (binary decisions: 1.0 = item1 wins, 0.0 = item2 wins)
            accumulation_steps: Number of steps to accumulate gradients over
            
        Returns:
            Loss value for this step
        """
        batch = {key: value.to(self.device) for key, value in batch.items()}
        
        with self._autocast_context():
            valence_A, ad_A = self.model(
                batch['input_ids_A'],
                batch['attention_mask_A']
            )
            valence_B, ad_B = self.model(
                batch['input_ids_B'],
                batch['attention_mask_B']
            )

            # Squeeze only the feature dimension to prevent batch_size=1 bugs
            pred_V_A = valence_A.squeeze(-1)
            pred_V_B = valence_B.squeeze(-1)
            
            # Valence: Pointwise MSE loss against gold scores (1-5 scale)
            loss_V = (F.mse_loss(pred_V_A, batch['V_A']) + \
                      F.mse_loss(pred_V_B, batch['V_B'])) / 2.0
            
            # Arousal & Dominance: Pairwise binary cross-entropy loss
            # loss = -[label * log(σ(r_A - r_B)) + (1-label) * log(σ(r_B - r_A))]
            # label = 1.0 if item1 wins, 0.0 if item2 wins
            # eps = 1e-7
            
            # Arousal (index 0 in arousal_dominance output)
            label_A = batch['label_A']
            logit_A = ad_A[:, 0] - ad_B[:, 0]
            loss_A = F.binary_cross_entropy_with_logits(logit_A, label_A)
            # per_sample_loss_A = -(label_A * torch.log(torch.sigmoid(logit_A) + eps) + 
            #            (1 - label_A) * torch.log(torch.sigmoid(-logit_A) + eps))
            # loss_A = per_sample_loss_A.mean()
            
            # Dominance (index 1 in arousal_dominance output)
            label_D = batch['label_D']
            logit_D = ad_A[:, 1] - ad_B[:, 1]
            loss_D = F.binary_cross_entropy_with_logits(logit_D, label_D)
            # per_sample_loss_D = -(label_D * torch.log(torch.sigmoid(logit_D) + eps) + 
            #            (1 - label_D) * torch.log(torch.sigmoid(-logit_D) + eps))
            # loss_D = per_sample_loss_D.mean()
            
            total_loss = (loss_V + loss_A + loss_D) / accumulation_steps
        
        self.scaler.scale(total_loss).backward()
        return total_loss.item() * accumulation_steps  # Return the unscaled loss for logging
    
    def evaluate(self, eval_loader: DataLoader) -> Dict[str, float]:
        """Evaluate the model: MSE for V (pointwise), Bradley-Terry for A/D (pairwise)."""
        self.model.eval()
        
        total_loss = 0
        losses_by_dim = {'V': 0, 'A': 0, 'D': 0}
        
        total_correct_V = 0
        total_correct_A = 0
        total_correct_D = 0
        total_count_V = 0
        total_count_A = 0
        total_count_D = 0
        
        # eps = 1e-7
        
        with torch.no_grad():
            for batch in tqdm(eval_loader, desc="Evaluating"):
                batch = {key: value.to(self.device) for key, value in batch.items()}
                
                valence_A, ad_A = self.model(
                    batch['input_ids_A'],
                    batch['attention_mask_A']
                )
                valence_B, ad_B = self.model(
                    batch['input_ids_B'],
                    batch['attention_mask_B']
                )

                # Squeeze only the feature dimension to prevent batch_size=1 bugs
                pred_V_A = valence_A.squeeze(-1)
                pred_V_B = valence_B.squeeze(-1)
                
                # Valence: MSE loss against gold scores
                loss_V = (F.mse_loss(pred_V_A, batch['V_A']) + \
                         F.mse_loss(pred_V_B, batch['V_B'])) / 2.0
                
                # Arousal: binary cross-entropy loss
                label_A = batch['label_A']
                logit_A = ad_A[:, 0] - ad_B[:, 0]
                loss_A = F.binary_cross_entropy_with_logits(logit_A, label_A)

                # per_sample_loss_A = -(label_A * torch.log(torch.sigmoid(logit_A) + eps) + 
                #           (1 - label_A) * torch.log(torch.sigmoid(-logit_A) + eps))
                
                # loss_A = per_sample_loss_A.mean()
                
                # Dominance: binary cross-entropy loss
                label_D = batch['label_D']
                logit_D = ad_A[:, 1] - ad_B[:, 1]
                loss_D = F.binary_cross_entropy_with_logits(logit_D, label_D)

                # per_sample_loss_D = -(label_D * torch.log(torch.sigmoid(logit_D) + eps) + 
                #           (1 - label_D) * torch.log(torch.sigmoid(-logit_D) + eps))
                # loss_D = per_sample_loss_D.mean()
                
                batch_loss = loss_V + loss_A + loss_D
                total_loss += batch_loss.item()
                losses_by_dim['V'] += loss_V.item()
                losses_by_dim['A'] += loss_A.item()
                losses_by_dim['D'] += loss_D.item()
                
                # Valence pairwise accuracy via gold EmoBank V scores
                V_A_gold = batch['V_A']
                V_B_gold = batch['V_B']
                correct_V = ((V_A_gold > V_B_gold) & (pred_V_A > pred_V_B)) | \
                           ((V_A_gold < V_B_gold) & (pred_V_A < pred_V_B))
                non_tie_V = (V_A_gold != V_B_gold)
                total_correct_V += correct_V[non_tie_V].sum().item()
                total_count_V += non_tie_V.sum().item()
                
                # Arousal/Dominance pairwise accuracy
                correct_A = ((label_A == 1.0) & (ad_A[:, 0] > ad_B[:, 0])) | \
                           ((label_A == 0.0) & (ad_B[:, 0] > ad_A[:, 0]))
                correct_D = ((label_D == 1.0) & (ad_A[:, 1] > ad_B[:, 1])) | \
                           ((label_D == 0.0) & (ad_B[:, 1] > ad_A[:, 1]))
                
                total_correct_A += correct_A.sum().item()
                total_correct_D += correct_D.sum().item()
                total_count_A += label_A.shape[0]
                total_count_D += label_D.shape[0]
        
        self.model.train()
        
        accuracy_V = total_correct_V / max(total_count_V, 1)
        accuracy_A = total_correct_A / max(total_count_A, 1)
        accuracy_D = total_correct_D / max(total_count_D, 1)
        accuracy_overall = (accuracy_V + accuracy_A + accuracy_D) / 3
        
        return {
            'eval_loss': total_loss / len(eval_loader),
            'eval_loss_V': losses_by_dim['V'] / len(eval_loader),
            'eval_loss_A': losses_by_dim['A'] / len(eval_loader),
            'eval_loss_D': losses_by_dim['D'] / len(eval_loader),
            'eval_accuracy': accuracy_overall,
            'eval_accuracy_V': accuracy_V,
            'eval_accuracy_A': accuracy_A,
            'eval_accuracy_D': accuracy_D
        }
    
    def score_text(self, text: str) -> Dict[str, float]:
        """
        Score a single text item across all three dimensions (V, A, D).
        
        Args:
            text: Text to score
            
        Returns:
            Dictionary with 'valence', 'arousal', 'dominance' scores
        """
        self.model.eval()
        
        encoding = self.tokenizer(
            text,
            max_length=self.max_length,
            padding='max_length',
            truncation=True,
            return_tensors='pt'
        )
        
        with torch.no_grad():
            valence, arousal_dominance = self.model(
                encoding['input_ids'].to(self.device),
                encoding['attention_mask'].to(self.device)
            )
        
        self.model.train()
        return {
            'valence': valence[0, 0].item(),
            'arousal': arousal_dominance[0, 0].item(),
            'dominance': arousal_dominance[0, 1].item()
        }
    
    def score_batch(self, texts: List[str], batch_size: int = 32) -> Dict[str, np.ndarray]:
        """
        Score multiple texts efficiently across all three dimensions.
        
        Args:
            texts: List of texts to score
            batch_size: Batch size for processing
            
        Returns:
            Dictionary with 'valence', 'arousal', 'dominance' arrays
        """
        self.model.eval()
        scores_V = []
        scores_A = []
        scores_D = []
        
        for i in range(0, len(texts), batch_size):
            batch_texts = texts[i:i + batch_size]
            encodings = self.tokenizer(
                batch_texts,
                max_length=self.max_length,
                padding='max_length',
                truncation=True,
                return_tensors='pt'
            )
            
            with torch.no_grad():
                valence, arousal_dominance = self.model(
                    encodings['input_ids'].to(self.device),
                    encodings['attention_mask'].to(self.device)
                )
            
            scores_V.extend(valence[:, 0].cpu().numpy())
            scores_A.extend(arousal_dominance[:, 0].cpu().numpy())
            scores_D.extend(arousal_dominance[:, 1].cpu().numpy())
        
        self.model.train()
        return {
            'valence': np.array(scores_V),
            'arousal': np.array(scores_A),
            'dominance': np.array(scores_D)
        }
    
    def normalize_scores(
        self,
        scores: np.ndarray,
        scale_min: float = 0.0,
        scale_max: float = 1.0
    ) -> np.ndarray:
        """
        Normalize raw reward scores to a desired scale.
        This normalizes within the provided scores (relative scaling).
        For consistent scaling across splits, use normalize_scores_with_params instead.
        
        Args:
            scores: Raw scores to normalize
            scale_min: Minimum value of output scale
            scale_max: Maximum value of output scale
            
        Returns:
            Normalized scores
        """
        score_min = scores.min()
        score_max = scores.max()
        
        if score_max == score_min:
            return np.full_like(scores, (scale_min + scale_max) / 2)
        
        normalized = (scores - score_min) / (score_max - score_min)
        normalized = normalized * (scale_max - scale_min) + scale_min
        
        return normalized
    
    def test_model(self, test_loader: DataLoader) -> Dict[str, float]:
        """Evaluate on the test set: MSE for V, Bradley-Terry for A/D."""
        print("\n" + "="*60)
        print("Running Test Evaluation (V, A, D)")
        print("="*60)
        
        self.model.eval()
        total_loss = 0
        losses_by_dim = {'V': 0, 'A': 0, 'D': 0}
        
        total_correct_V = 0
        total_correct_A = 0
        total_correct_D = 0
        total_count_V = 0
        total_count_A = 0
        total_count_D = 0
        
        eps = 1e-7
        
        with torch.no_grad():
            for batch in tqdm(test_loader, desc="Testing"):
                batch = {key: value.to(self.device) for key, value in batch.items()}
                
                valence_A, ad_A = self.model(
                    batch['input_ids_A'],
                    batch['attention_mask_A']
                )
                valence_B, ad_B = self.model(
                    batch['input_ids_B'],
                    batch['attention_mask_B']
                )

                # Squeeze only the feature dimension to prevent batch_size=1 bugs
                pred_V_A = valence_A.squeeze(-1)
                pred_V_B = valence_B.squeeze(-1)
                
                # Valence: MSE loss
                loss_V = (F.mse_loss(pred_V_A, batch['V_A']) + \
                          F.mse_loss(pred_V_B, batch['V_B'])) / 2.0
                
                # Arousal: binary cross-entropy loss
                label_A = batch['label_A']
                logit_A = ad_A[:, 0] - ad_B[:, 0]
                loss_A = F.binary_cross_entropy_with_logits(logit_A, label_A)
                # per_sample_loss_A = -(label_A * torch.log(torch.sigmoid(logit_A) + eps) + 
                #           (1 - label_A) * torch.log(torch.sigmoid(-logit_A) + eps))
                # loss_A = per_sample_loss_A.mean()
                
                # Dominance: binary cross-entropy loss
                label_D = batch['label_D']
                logit_D = ad_A[:, 1] - ad_B[:, 1]
                loss_D = F.binary_cross_entropy_with_logits(logit_D, label_D)
                # per_sample_loss_D = -(label_D * torch.log(torch.sigmoid(logit_D) + eps) + 
                #           (1 - label_D) * torch.log(torch.sigmoid(-logit_D) + eps))
                # loss_D = per_sample_loss_D.mean()
                
                batch_loss = loss_V + loss_A + loss_D
                total_loss += batch_loss.item()
                losses_by_dim['V'] += loss_V.item()
                losses_by_dim['A'] += loss_A.item()
                losses_by_dim['D'] += loss_D.item()
                
                # Valence pairwise accuracy via gold EmoBank V scores
                V_A_gold = batch['V_A']
                V_B_gold = batch['V_B']
                correct_V = ((V_A_gold > V_B_gold) & (pred_V_A > pred_V_B)) | \
                           ((V_A_gold < V_B_gold) & (pred_V_A < pred_V_B))
                non_tie_V = (V_A_gold != V_B_gold)
                total_correct_V += correct_V[non_tie_V].sum().item()
                total_count_V += non_tie_V.sum().item()
                
                # Arousal/Dominance pairwise accuracy
                correct_A = ((label_A == 1.0) & (ad_A[:, 0] > ad_B[:, 0])) | \
                           ((label_A == 0.0) & (ad_B[:, 0] > ad_A[:, 0]))
                correct_D = ((label_D == 1.0) & (ad_A[:, 1] > ad_B[:, 1])) | \
                           ((label_D == 0.0) & (ad_B[:, 1] > ad_A[:, 1]))
                
                total_correct_A += correct_A.sum().item()
                total_correct_D += correct_D.sum().item()
                total_count_A += label_A.shape[0]
                total_count_D += label_D.shape[0]
        
        test_loss = total_loss / len(test_loader)
        test_loss_V = losses_by_dim['V'] / len(test_loader)
        test_loss_A = losses_by_dim['A'] / len(test_loader)
        test_loss_D = losses_by_dim['D'] / len(test_loader)
        
        accuracy_V = total_correct_V / max(total_count_V, 1)
        accuracy_A = total_correct_A / max(total_count_A, 1)
        accuracy_D = total_correct_D / max(total_count_D, 1)
        accuracy_overall = (accuracy_V + accuracy_A + accuracy_D) / 3
        
        print(f"\nTest Results:")
        print(f"  Total Loss: {test_loss:.4f} | Accuracy: {accuracy_overall:.2%}")
        print(f"  Loss - V (MSE): {test_loss_V:.4f} | A: {test_loss_A:.4f} | D: {test_loss_D:.4f}")
        print(f"  Acc  - V (EMOBANK): {accuracy_V:.2%} | A: {accuracy_A:.2%} | D: {accuracy_D:.2%}")
        print("\n" + "="*60 + "\n")
        
        results = {
            'test_loss': test_loss,
            'test_loss_V': test_loss_V,
            'test_loss_A': test_loss_A,
            'test_loss_D': test_loss_D,
            'test_accuracy': accuracy_overall,
            'test_accuracy_V': accuracy_V,
            'test_accuracy_A': accuracy_A,
            'test_accuracy_D': accuracy_D
        }
        
        return results
    
    def save(self, path: str):
        """Save model and training state."""
        torch.save({
            'model_state_dict': self.model.state_dict(),
            'optimizer_state_dict': self.optimizer.state_dict() if self.optimizer else None,
            'scheduler_state_dict': self.scheduler.state_dict() if self.scheduler else None,
            'training_history': self.training_history,
            'config': {
                'model_name': self.model_name,
                'max_length': self.max_length
            }
        }, path)
        print(f"Model saved to {path}")
    
    def load(self, path: str):
        """Load model and training state."""
        checkpoint = torch.load(path, map_location=self.device)
        self.model.load_state_dict(checkpoint['model_state_dict'])
        
        if checkpoint['optimizer_state_dict'] and self.optimizer:
            self.optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
        
        if checkpoint['scheduler_state_dict'] and self.scheduler:
            self.scheduler.load_state_dict(checkpoint['scheduler_state_dict'])
        
        self.training_history = checkpoint.get('training_history', [])
        print(f"Model loaded from {path}")

    def push_to_hub(self, repo_id: str, private: bool = True):
        """
        Push the trained model to HuggingFace Hub.
        
        Args:
            repo_id: Repository ID in format "username/repo-name"
            private: Whether to make the repository private (default True)
        """
        from huggingface_hub import create_repo, upload_folder
        import os
        
        # Create a temporary directory for the model files
        temp_dir = f"./temp_model_{repo_id.split('/')[-1]}"
        os.makedirs(temp_dir, exist_ok=True)
        
        try:
            # Save encoder and reward head
            self.model.encoder.save_pretrained(os.path.join(temp_dir, "encoder"))
            self.tokenizer.save_pretrained(temp_dir)
            torch.save(
                {
                    "model_state_dict": self.model.state_dict(),
                    "config": {
                        "model_name": self.model_name,
                        "max_length": self.max_length,
                    },
                },
                os.path.join(temp_dir, "reward_model.pth"),
            )
            
            # Save metadata
            metadata = {
                'model_name': self.model_name,
                'max_length': self.max_length,
                'training_history': self.training_history
            }
            with open(os.path.join(temp_dir, "training_metadata.json"), 'w') as f:
                json.dump(metadata, f, indent=2)
            
            # Create and upload to Hub
            create_repo(repo_id, private=private, exist_ok=True)
            upload_folder(repo_name=repo_id, folder_path=temp_dir, repo_type="model")
            
            print(f"Model successfully pushed to HuggingFace Hub: https://huggingface.co/{repo_id}")
        
        finally:
            # Clean up temporary directory
            import shutil
            if os.path.exists(temp_dir):
                shutil.rmtree(temp_dir)