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)