fix: silly mistake (#82)

* fix: a very stupid bug * chore: refac
tatsu-lab · Dec 3, 2023 · f8b5fb1 · f8b5fb1
1 parent 6f53082
commit f8b5fb1
Showing 1 changed file with 5 additions and 4 deletions.
diff --git a/src/alpaca_farm/rl/dpo_trainer.py b/src/alpaca_farm/rl/dpo_trainer.py
@@ -38,14 +38,15 @@ def compute_loss(self, model, inputs, return_outputs=False):
  with torch.no_grad():
  ref_logits_w = self.ref_model(input_ids=input_ids_w, attention_mask=attention_mask_w).logits[..., :-1, :]
  ref_logits_l = self.ref_model(input_ids=input_ids_l, attention_mask=attention_mask_l).logits[..., :-1, :]
- ref_logprobs_w = F.cross_entropy(ref_logits_w.transpose(-1, -2), labels_w, reduction="none").sum(-1)
- ref_logprobs_l = F.cross_entropy(ref_logits_l.transpose(-1, -2), labels_l, reduction="none").sum(-1)
+ ref_logprobs_w = -F.cross_entropy(ref_logits_w.transpose(-1, -2), labels_w, reduction="none").sum(-1)
+ ref_logprobs_l = -F.cross_entropy(ref_logits_l.transpose(-1, -2), labels_l, reduction="none").sum(-1)
 
  logits_w = model(input_ids=input_ids_w, attention_mask=attention_mask_w).logits[..., :-1, :]
  logits_l = model(input_ids=input_ids_l, attention_mask=attention_mask_l).logits[..., :-1, :]
- logprobs_w = F.cross_entropy(logits_w.transpose(-1, -2), labels_w, reduction="none").sum(-1)
- logprobs_l = F.cross_entropy(logits_l.transpose(-1, -2), labels_l, reduction="none").sum(-1)
+ logprobs_w = -F.cross_entropy(logits_w.transpose(-1, -2), labels_w, reduction="none").sum(-1)
+ logprobs_l = -F.cross_entropy(logits_l.transpose(-1, -2), labels_l, reduction="none").sum(-1)
 
  logits = self.args.beta * ((logprobs_w - ref_logprobs_w) - (logprobs_l - ref_logprobs_l))
  loss = -F.logsigmoid(logits).mean(0)
+
  return (loss, dict(logits=logits)) if return_outputs else loss