cehrgpt 0.0.2__py3-none-any.whl → 0.1.1__py3-none-any.whl

cehrgpt/rl_finetune/cehrgpt_dpo_trainer.py

@@ -1,586 +0,0 @@
-import warnings
-from collections import defaultdict
-from copy import deepcopy
-from typing import Any, Callable, Dict, List, Literal, Optional, Tuple, Union
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from accelerate.utils import is_deepspeed_available
-from datasets import Dataset
-from transformers import PreTrainedModel, Trainer
-from transformers.trainer_callback import TrainerCallback
-from transformers.trainer_utils import EvalLoopOutput
-from trl.trainer.callbacks import SyncRefModelCallback
-from trl.trainer.dpo_config import DPOConfig, FDivergenceConstants, FDivergenceType
-from trl.trainer.utils import RunningMoments, cap_exp, disable_dropout_in_model
-
-from cehrgpt.data.hf_cehrgpt_dpo_collator import CehrGptDPODataCollator
-from cehrgpt.models.hf_cehrgpt import CEHRGPT2LMHeadModel
-from cehrgpt.models.tokenization_hf_cehrgpt import CehrGptTokenizer
-
-if is_deepspeed_available():
-    import deepspeed
-
-
-class CehrGptDPOTrainer(Trainer):
-    def __init__(
-        self,
-        model: CEHRGPT2LMHeadModel,
-        ref_model: CEHRGPT2LMHeadModel,
-        tokenizer: CehrGptTokenizer,
-        args: DPOConfig,
-        data_collator: CehrGptDPODataCollator,
-        train_dataset: Optional[Dataset] = None,
-        eval_dataset: Optional[Union[Dataset, Dict[str, Dataset]]] = None,
-        callbacks: Optional[List[TrainerCallback]] = None,
-        optimizers: Tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (
-            None,
-            None,
-        ),
-        preprocess_logits_for_metrics: Optional[
-            Callable[[torch.Tensor, torch.Tensor], torch.Tensor]
-        ] = None,
-        disable_dropout: bool = True,
-        compute_metrics: Optional[Callable[[EvalLoopOutput], Dict]] = None,
-    ):
-        if ref_model is model:
-            raise ValueError(
-                "`model` and `ref_model` cannot be the same object. If you want `ref_model` to be the "
-                "same as `model`, you must mass a copy of it, or `None` if you use peft."
-            )
-
-        if getattr(args, "gradient_checkpointing", False):
-            # For backward compatibility with older versions of transformers
-            if hasattr(model, "enable_input_require_grads"):
-                model.enable_input_require_grads()
-            else:
-
-                def make_inputs_require_grad(module, input, output):
-                    output.requires_grad_(True)
-
-                model.get_input_embeddings().register_forward_hook(
-                    make_inputs_require_grad
-                )
-
-        if model is not None:
-            self.is_encoder_decoder = model.config.is_encoder_decoder
-        elif args.is_encoder_decoder is None:
-            raise ValueError(
-                "When no model is provided, you need to pass the parameter is_encoder_decoder to the DPOTrainer/DPOConfig."
-            )
-        else:
-            self.is_encoder_decoder = args.is_encoder_decoder
-
-        self.tokenizer = tokenizer
-        self.ref_model = ref_model
-
-        if not disable_dropout:
-            warnings.warn(
-                "You passed `disable_dropout` to the DPOTrainer, the value you passed will override the one in the `DPOConfig`."
-            )
-            args.disable_dropout = disable_dropout
-
-        if args.disable_dropout:
-            disable_dropout_in_model(model)
-            if self.ref_model is not None:
-                disable_dropout_in_model(self.ref_model)
-
-        self.beta = args.beta
-        self.label_smoothing = args.label_smoothing
-        self.loss_type = args.loss_type
-        self._stored_metrics = defaultdict(lambda: defaultdict(list))
-        self.f_divergence_type = args.f_divergence_type
-        self.f_divergence_params = {
-            FDivergenceConstants.ALPHA_DIVERGENCE_COEF_KEY: args.f_alpha_divergence_coef
-        }
-        self.reference_free = args.reference_free
-        # α parameter from the [RPO](https://huggingface.co/papers/2404.19733) paper (v3), which controls the
-        # weighting of the NLL term in the loss. If `None`, no weighting is applied and the loss is the same as the
-        # DPO loss. The paper recommends `rpo_alpha=1.0`.
-        self.rpo_alpha = args.rpo_alpha
-        self.vs_token_id = tokenizer._convert_token_to_id("VS")
-        if self.vs_token_id == tokenizer._oov_token_id:
-            self.vs_token_id = tokenizer._convert_token_to_id("[VS]")
-        self.ve_token_id = tokenizer._convert_token_to_id("VE")
-        if self.ve_token_id == tokenizer._oov_token_id:
-            self.ve_token_id = tokenizer._convert_token_to_id("[VE]")
-        self.non_active_token_ids = [
-            tokenizer.pad_token_id,
-            self.vs_token_id,
-            self.ve_token_id,
-        ]
-
-        super().__init__(
-            model=model,
-            args=args,
-            data_collator=data_collator,
-            train_dataset=train_dataset,
-            eval_dataset=eval_dataset,
-            tokenizer=tokenizer,
-            compute_metrics=compute_metrics,
-            callbacks=callbacks,
-            optimizers=optimizers,
-            preprocess_logits_for_metrics=preprocess_logits_for_metrics,
-        )
-
-        if not hasattr(self, "accelerator"):
-            raise AttributeError(
-                "Your `Trainer` does not have an `accelerator` object. Consider upgrading `transformers`."
-            )
-
-        if self.is_deepspeed_enabled:
-            self.ref_model = self._prepare_deepspeed(self.ref_model)
-        else:
-            self.ref_model = self.accelerator.prepare_model(
-                self.ref_model, evaluation_mode=True
-            )
-            self.add_callback(
-                SyncRefModelCallback(
-                    ref_model=self.ref_model, accelerator=self.accelerator
-                )
-            )
-        if self.loss_type == "bco_pair":
-            self.running = RunningMoments(self.accelerator)
-
-    def _prepare_deepspeed(self, model: CEHRGPT2LMHeadModel):
-        # Adapted from accelerate: https://github.com/huggingface/accelerate/blob/739b135f8367becb67ffaada12fe76e3aa60fefd/src/accelerate/accelerator.py#L1473
-        deepspeed_plugin = self.accelerator.state.deepspeed_plugin
-        config_kwargs = deepcopy(deepspeed_plugin.deepspeed_config)
-
-        if model is not None:
-            if hasattr(model, "config"):
-                hidden_size = (
-                    max(model.config.hidden_sizes)
-                    if getattr(model.config, "hidden_sizes", None)
-                    else getattr(model.config, "hidden_size", None)
-                )
-                if (
-                    hidden_size is not None
-                    and config_kwargs["zero_optimization"]["stage"] == 3
-                ):
-                    # Note that `stage3_prefetch_bucket_size` can produce DeepSpeed messages like: `Invalidate trace cache @ step 0: expected module 1, but got module 0`
-                    # This is expected and is not an error, see: https://github.com/microsoft/DeepSpeed/discussions/4081
-                    config_kwargs.update(
-                        {
-                            "zero_optimization.reduce_bucket_size": hidden_size
-                            * hidden_size,
-                            "zero_optimization.stage3_param_persistence_threshold": 10
-                            * hidden_size,
-                            "zero_optimization.stage3_prefetch_bucket_size": 0.9
-                            * hidden_size
-                            * hidden_size,
-                        }
-                    )
-
-        # If ZeRO-3 is used, we shard both the active and reference model.
-        # Otherwise, we assume the reference model fits in memory and is initialized on each device with ZeRO disabled (stage 0)
-        if config_kwargs["zero_optimization"]["stage"] != 3:
-            config_kwargs["zero_optimization"]["stage"] = 0
-        model, *_ = deepspeed.initialize(model=model, config=config_kwargs)
-        model.eval()
-        return model
-
-    def forward(
-        self,
-        model: CEHRGPT2LMHeadModel,
-        batch: Dict[str, torch.Tensor],
-    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
-        """Run the given model on the given batch of inputs, concatenating the chosen and rejected inputs together.
-
-        We do this to avoid doing two forward passes, because it's faster for FSDP.
-        """
-        chosen_outputs = model(
-            input_ids=batch["chosen_input_ids"],
-            attention_mask=batch["chosen_attention_mask"],
-            value_indicators=(
-                batch["chosen_value_indicators"]
-                if "chosen_value_indicators" in batch
-                else None
-            ),
-            values=batch["chosen_values"] if "chosen_values" in batch else None,
-        )
-
-        chosen_logps, chosen_logits = self.get_batch_logps(
-            chosen_outputs.logits,
-            batch["chosen_input_ids"],
-            pad_token_ids=[self.tokenizer.pad_token_id],
-        )
-
-        # move labels to correct device to enable model parallelism
-        labels = batch["chosen_input_ids"].to(chosen_outputs.logits.device)
-        # Shift so that tokens < n predict n
-        shift_logits = chosen_outputs.logits[..., :-1, :].contiguous()
-        shift_labels = labels[..., 1:].contiguous()
-        # Flatten the tokens
-        loss_fct = nn.CrossEntropyLoss()
-        nll_loss = loss_fct(
-            shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1)
-        )
-
-        rejected_outputs = model(
-            input_ids=batch["rejected_input_ids"],
-            attention_mask=batch["rejected_attention_mask"],
-            value_indicators=(
-                batch["rejected_value_indicators"]
-                if "rejected_value_indicators" in batch
-                else None
-            ),
-            values=batch["rejected_values"] if "rejected_values" in batch else None,
-        )
-
-        rejected_logps, rejected_logits = self.get_batch_logps(
-            rejected_outputs.logits,
-            batch["rejected_input_ids"],
-            pad_token_ids=[self.tokenizer.pad_token_id],
-        )
-
-        return (
-            chosen_logps,
-            rejected_logps,
-            chosen_logits,
-            rejected_logits,
-            nll_loss,
-        )
-
-    def dpo_loss(
-        self,
-        policy_chosen_logps: torch.Tensor,
-        policy_rejected_logps: torch.Tensor,
-        reference_chosen_logps: torch.Tensor,
-        reference_rejected_logps: torch.Tensor,
-    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-        """Compute the DPO loss for a batch of policy and reference model log probabilities.
-
-        Args:
-            policy_chosen_logps: Log probabilities of the policy model for the chosen responses. Shape: (batch_size,)
-            policy_rejected_logps: Log probabilities of the policy model for the rejected responses. Shape: (batch_size,)
-            reference_chosen_logps: Log probabilities of the reference model for the chosen responses. Shape: (batch_size,)
-            reference_rejected_logps: Log probabilities of the reference model for the rejected responses. Shape: (batch_size,)
-
-        Returns:
-            A tuple of three tensors: (losses, chosen_rewards, rejected_rewards).
-            The losses tensor contains the DPO loss for each example in the batch.
-            The chosen_rewards and rejected_rewards tensors contain the rewards for the chosen and rejected responses, respectively.
-        """
-        chosen_logratios = policy_chosen_logps.to(self.accelerator.device) - (
-            not self.reference_free
-        ) * reference_chosen_logps.to(self.accelerator.device)
-        rejected_logratios = policy_rejected_logps.to(self.accelerator.device) - (
-            not self.reference_free
-        ) * reference_rejected_logps.to(self.accelerator.device)
-
-        if self.f_divergence_type == FDivergenceType.ALPHA_DIVERGENCE.value:
-            # The alpha-divergence formula: (1 - u^-alpha) / alpha
-            # The divergence difference between the chosen and rejected sample is:
-            #     (1 - u[w]^-alpha) / alpha - (1 - u[l]^-alpha) / alpha
-            #        = (u[l]^-alpha - u[w]^-alpha) / alpha
-            # where u[w] and u[l] are the policy/reference probability ratios
-            # for the chosen and rejected samples, respectively.
-            alpha_coef = FDivergenceConstants.ALPHA_DIVERGENCE_COEF_DEFAULT
-            if (
-                self.f_divergence_params
-                and FDivergenceConstants.ALPHA_DIVERGENCE_COEF_KEY
-                in self.f_divergence_params
-            ):
-                alpha_coef = float(
-                    self.f_divergence_params[
-                        FDivergenceConstants.ALPHA_DIVERGENCE_COEF_KEY
-                    ]
-                )
-            logits = (
-                cap_exp(rejected_logratios * -alpha_coef)
-                - cap_exp(chosen_logratios * -alpha_coef)
-            ) / alpha_coef
-        else:
-            pi_logratios = policy_chosen_logps - policy_rejected_logps
-            if self.reference_free:
-                ref_logratios = torch.tensor(
-                    [0], dtype=pi_logratios.dtype, device=pi_logratios.device
-                )
-            else:
-                ref_logratios = reference_chosen_logps - reference_rejected_logps
-
-            pi_logratios = pi_logratios.to(self.accelerator.device)
-            ref_logratios = ref_logratios.to(self.accelerator.device)
-            logits = pi_logratios - ref_logratios
-
-            if self.f_divergence_type == FDivergenceType.JS_DIVERGENCE.value:
-                # The js-divergence formula: log(2 * u / (1 + u))
-                # The divergence difference between the chosen and rejected sample is:
-                #     log(2 * u[w] / (1 + u[w])) - log(2 * u[l] / (1 + u[l]))
-                #       = log(u[w]) - log(u[l]) - (log(1 + u[w]) - log(1 + u[l]))
-                # where u[w] and u[l] are the policy/reference probability ratios
-                # for the chosen and rejected samples, respectively.
-                logits -= F.softplus(chosen_logratios) - F.softplus(rejected_logratios)
-
-        # The beta is a temperature parameter for the DPO loss, typically something in the range of 0.1 to 0.5.
-        # We ignore the reference model as beta -> 0. The label_smoothing parameter encodes our uncertainty about the labels and
-        # calculates a conservative DPO loss.
-        if self.loss_type == "sigmoid":
-            losses = (
-                -F.logsigmoid(self.beta * logits) * (1 - self.label_smoothing)
-                - F.logsigmoid(-self.beta * logits) * self.label_smoothing
-            )
-        elif self.loss_type == "robust":
-            losses = (
-                -F.logsigmoid(self.beta * logits) * (1 - self.label_smoothing)
-                + F.logsigmoid(-self.beta * logits) * self.label_smoothing
-            ) / (1 - 2 * self.label_smoothing)
-        elif self.loss_type == "exo_pair":
-            # eqn (16) of the EXO paper: https://huggingface.co/papers/2402.00856
-            import math
-
-            if self.label_smoothing == 0:
-                self.label_smoothing = 1e-3
-            losses = (self.beta * logits).sigmoid() * (
-                F.logsigmoid(self.beta * logits) - math.log(1 - self.label_smoothing)
-            ) + (-self.beta * logits).sigmoid() * (
-                F.logsigmoid(-self.beta * logits) - math.log(self.label_smoothing)
-            )
-        elif self.loss_type == "hinge":
-            losses = torch.relu(1 - self.beta * logits)
-        elif self.loss_type == "ipo":
-            # eqn (17) of the paper where beta is the regularization parameter for the IPO loss, denoted by tau in the paper.
-            losses = (logits - 1 / (2 * self.beta)) ** 2
-        elif self.loss_type == "bco_pair":
-            chosen_logratios = policy_chosen_logps - reference_chosen_logps
-            rejected_logratios = policy_rejected_logps - reference_rejected_logps
-
-            chosen_rewards = self.beta * chosen_logratios
-            rejected_rewards = self.beta * rejected_logratios
-            rewards = torch.cat((chosen_rewards, rejected_rewards), 0).mean().detach()
-            self.running.update(rewards)
-            delta = self.running.mean
-
-            losses = -F.logsigmoid(
-                (self.beta * chosen_logratios) - delta
-            ) - F.logsigmoid(-(self.beta * rejected_logratios - delta))
-        elif self.loss_type == "sppo_hard":
-            # In the paper (https://huggingface.co/papers/2405.00675), SPPO employs a soft probability approach, estimated using the PairRM score. The probability calculation is conducted outside of the trainer class. The version described here is the hard probability version, where P in Equation (4.7) of Algorithm 1 is set to 1 for the winner and 0 for the loser.
-            a = policy_chosen_logps - reference_chosen_logps
-            b = policy_rejected_logps - reference_rejected_logps
-
-            losses = (a - 0.5 / self.beta) ** 2 + (b + 0.5 / self.beta) ** 2
-        elif self.loss_type == "nca_pair":
-            chosen_rewards = (policy_chosen_logps - reference_chosen_logps) * self.beta
-            rejected_rewards = (
-                policy_rejected_logps - reference_rejected_logps
-            ) * self.beta
-            losses = (
-                -F.logsigmoid(chosen_rewards)
-                - 0.5 * F.logsigmoid(-chosen_rewards)
-                - 0.5 * F.logsigmoid(-rejected_rewards)
-            )
-        elif self.loss_type == "aot_pair":
-            chosen_logratios = policy_chosen_logps - reference_chosen_logps
-            rejected_logratios = policy_rejected_logps - reference_rejected_logps
-
-            chosen_logratios_sorted, _ = torch.sort(chosen_logratios, dim=0)
-            rejected_logratios_sorted, _ = torch.sort(rejected_logratios, dim=0)
-
-            delta = chosen_logratios_sorted - rejected_logratios_sorted
-
-            losses = (
-                -F.logsigmoid(self.beta * delta) * (1 - self.label_smoothing)
-                - F.logsigmoid(-self.beta * delta) * self.label_smoothing
-            )
-
-        elif self.loss_type == "aot":
-            pi_logratios = policy_chosen_logps - policy_rejected_logps
-            ref_logratios = reference_chosen_logps - reference_rejected_logps
-
-            pi_logratios_sorted, _ = torch.sort(pi_logratios, dim=0)
-            ref_logratios_sorted, _ = torch.sort(ref_logratios, dim=0)
-
-            delta = pi_logratios_sorted - ref_logratios_sorted
-
-            losses = (
-                -F.logsigmoid(self.beta * delta) * (1 - self.label_smoothing)
-                - F.logsigmoid(-self.beta * delta) * self.label_smoothing
-            )
-
-        elif self.loss_type == "apo_zero":
-            # Eqn (7) of the APO paper (https://huggingface.co/papers/2408.06266)
-            # Use this loss when you believe the chosen outputs are better than your model's default output
-
-            losses_chosen = 1 - F.sigmoid(
-                self.beta * chosen_logratios
-            )  # Increase chosen likelihood
-            losses_rejected = F.sigmoid(
-                self.beta * rejected_logratios
-            )  # Decrease rejected likelihood
-
-            losses = losses_chosen + losses_rejected
-
-        elif self.loss_type == "apo_down":
-            # Eqn (8) of the APO paper (https://huggingface.co/papers/2408.06266)
-            # Use this loss when you believe the chosen outputs are worse than your model's default output
-
-            losses_chosen = F.sigmoid(
-                self.beta * chosen_logratios
-            )  # Decrease chosen likelihood
-            losses_rejected = 1 - F.sigmoid(
-                self.beta * (chosen_logratios - rejected_logratios)
-            )  # Decrease rejected likelihood more
-
-            losses = losses_chosen + losses_rejected
-
-        else:
-            raise ValueError(
-                f"Unknown loss type: {self.loss_type}. Should be one of ['sigmoid', 'hinge', 'ipo', 'exo_pair', 'nca_pair', 'robust', 'bco_pair', 'sppo_hard', 'aot', 'aot_pair', 'apo_zero', 'apo_down']"
-            )
-
-        chosen_rewards = (
-            self.beta
-            * (
-                policy_chosen_logps.to(self.accelerator.device)
-                - reference_chosen_logps.to(self.accelerator.device)
-            ).detach()
-        )
-        rejected_rewards = (
-            self.beta
-            * (
-                policy_rejected_logps.to(self.accelerator.device)
-                - reference_rejected_logps.to(self.accelerator.device)
-            ).detach()
-        )
-
-        return losses, chosen_rewards, rejected_rewards
-
-    def get_batch_loss_metrics(
-        self,
-        model,
-        batch: Dict[str, Union[List, torch.LongTensor]],
-        train_eval: Literal["train", "eval"] = "train",
-    ):
-        """Compute the DPO loss and other metrics for the given batch of inputs for train or test."""
-        metrics = {}
-        forward_output = self.forward(model, batch)
-        (
-            policy_chosen_logps,
-            policy_rejected_logps,
-            policy_chosen_logits,
-            policy_rejected_logits,
-            policy_nll_loss,
-        ) = forward_output[:5]
-
-        with torch.no_grad():
-            reference_chosen_logps, reference_rejected_logps = self.forward(
-                self.ref_model, batch
-            )[:2]
-
-        losses, chosen_rewards, rejected_rewards = self.dpo_loss(
-            policy_chosen_logps,
-            policy_rejected_logps,
-            reference_chosen_logps,
-            reference_rejected_logps,
-        )
-        reward_accuracies = (chosen_rewards > rejected_rewards).float()
-
-        if self.rpo_alpha is not None:
-            # RPO loss from V3 of the paper:
-            losses = losses + policy_nll_loss * self.rpo_alpha
-
-        prefix = "eval_" if train_eval == "eval" else ""
-        metrics[f"{prefix}rewards/chosen"] = chosen_rewards.mean().cpu()
-        metrics[f"{prefix}rewards/rejected"] = rejected_rewards.mean().cpu()
-        metrics[f"{prefix}rewards/accuracies"] = reward_accuracies.mean().cpu()
-        metrics[f"{prefix}rewards/margins"] = (
-            (chosen_rewards - rejected_rewards).mean().cpu()
-        )
-        metrics[f"{prefix}logps/rejected"] = policy_rejected_logps.detach().mean().cpu()
-        metrics[f"{prefix}logps/chosen"] = policy_chosen_logps.detach().mean().cpu()
-        metrics[f"{prefix}logits/rejected"] = (
-            policy_rejected_logits.detach().mean().cpu()
-        )
-        metrics[f"{prefix}logits/chosen"] = policy_chosen_logits.detach().mean().cpu()
-        if self.rpo_alpha is not None:
-            metrics[f"{prefix}nll_loss"] = policy_nll_loss.detach().mean().cpu()
-        return losses.mean(), metrics
-
-    @staticmethod
-    def get_batch_logps(
-        all_logits: torch.FloatTensor, labels: torch.Tensor, pad_token_ids: List[int]
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
-        # remove the demographics
-        all_logits = all_logits[:, 3:-1]
-        labels = labels.clone()[:, 4:]
-        # Calculate the sequence log probability log p(x)
-        per_token_logps = torch.gather(
-            all_logits.log_softmax(-1), dim=2, index=labels.unsqueeze(2)
-        ).squeeze(2)
-        non_pad_mask = ~torch.isin(
-            labels, torch.tensor(pad_token_ids, device=labels.device)
-        )
-        per_token_logps *= non_pad_mask
-        all_logps = per_token_logps.sum(-1)
-        return all_logps, all_logits
-
-    def compute_loss(
-        self,
-        model: Union[PreTrainedModel, nn.Module],
-        inputs: Dict[str, Union[torch.Tensor, Any]],
-        return_outputs=False,
-    ) -> Union[torch.Tensor, Tuple[torch.Tensor, Dict[str, torch.Tensor]]]:
-        loss, metrics = self.get_batch_loss_metrics(model, inputs, train_eval="train")
-        # Make sure to move the loss to the device the original accumulating loss is at back in the `Trainer` class:
-        loss = loss.to(self.args.device)
-        # force log the metrics
-        self.store_metrics(metrics, train_eval="train")
-        if return_outputs:
-            return (loss, metrics)
-        return loss
-
-    def store_metrics(
-        self, metrics: Dict[str, float], train_eval: Literal["train", "eval"] = "train"
-    ) -> None:
-        for key, value in metrics.items():
-            self._stored_metrics[train_eval][key].append(value)
-
-    def log(self, logs: Dict[str, float]) -> None:
-        """
-        Log `logs` on the various objects watching training, including stored metrics.
-
-        Args:
-            logs (`Dict[str, float]`):
-                The values to log.
-        """
-        # logs either has 'loss' or 'eval_loss'
-        train_eval = "train" if "loss" in logs else "eval"
-        # Add averaged stored metrics to logs
-        for key, metrics in self._stored_metrics[train_eval].items():
-            logs[key] = torch.tensor(metrics).mean().item()
-        del self._stored_metrics[train_eval]
-        return super().log(logs)
-
-    def prediction_step(
-        self,
-        model: Union[PreTrainedModel, nn.Module],
-        inputs: Dict[str, Union[torch.Tensor, Any]],
-        prediction_loss_only: bool,
-        ignore_keys: Optional[List[str]] = None,
-    ):
-        with torch.no_grad():
-            loss, metrics = self.get_batch_loss_metrics(
-                model, inputs, train_eval="eval"
-            )
-
-        # force log the metrics
-        self.store_metrics(metrics, train_eval="eval")
-
-        if prediction_loss_only:
-            return (loss.detach(), None, None)
-
-        # logits for the chosen and rejected samples from model
-        logits_dict = {
-            "eval_logits/chosen": metrics["eval_logits/chosen"],
-            "eval_logits/rejected": metrics["eval_logits/rejected"],
-        }
-        logits = tuple(
-            v.unsqueeze(dim=0) for k, v in logits_dict.items() if k not in ignore_keys
-        )
-        logits = torch.stack(logits).mean(axis=1).to(self.accelerator.device)
-        labels = torch.zeros(logits.shape[0], device=self.accelerator.device)
-
-        return (loss.detach(), logits, labels)