diff --git a/src/listings/ch07/bonus.rs b/src/listings/ch07/bonus.rs index b148cd4..edbbf0d 100644 --- a/src/listings/ch07/bonus.rs +++ b/src/listings/ch07/bonus.rs @@ -594,6 +594,32 @@ impl + Clone> Preference } } +#[allow(unused_variables)] +pub fn compute_dpo_loss( + model_chosen_logprobs: &Tensor, + model_rejected_logprobs: &Tensor, + reference_chosen_logprobs: &Tensor, + reference_rejected_logprobs: &Tensor, + beta: f64, +) -> Result<(Tensor, Tensor, Tensor)> { + let model_logratios = (model_chosen_logprobs - model_rejected_logprobs)?; + let reference_logratios = (reference_chosen_logprobs - reference_rejected_logprobs)?; + let logits = (model_logratios - reference_logratios)?; + + let mut losses = candle_nn::ops::sigmoid(&logits)?; + losses = (-1_f64 * losses.log()?)?; + + // Optional values to track progress during training + let chosen_rewards = (model_chosen_logprobs - reference_chosen_logprobs)?.detach(); + let rejected_rewards = (model_rejected_logprobs - reference_rejected_logprobs)?.detach(); + + Ok(( + losses.mean_all()?, + chosen_rewards.mean_all()?, + rejected_rewards.mean_all()?, + )) +} + #[cfg(test)] mod tests { use super::*;