RLHF take-home aligning a small open model

"""The core of the DPO step, without a framework, so it can be unit-tested.
Fill in the three functions; the training script calls them per batch."""


def sequence_logp(token_logps, loss_mask):
    """Sum of log-probabilities over the response tokens of one sequence.
    loss_mask is 1 on response tokens and 0 on prompt and padding tokens."""
    raise NotImplementedError


def dpo_loss_and_grad(policy_chosen, policy_rejected, ref_chosen, ref_rejected, beta=0.1):
    """DPO loss for one preference pair from sequence log-probabilities.
    Returns (loss, d_loss/d_policy_chosen, d_loss/d_policy_rejected).
    Must stay finite for very large margins."""
    raise NotImplementedError


def dpo_metrics(pairs, beta=0.1):
    """pairs: [(policy_chosen, policy_rejected, ref_chosen, ref_rejected)].
    Returns {"loss", "accuracy", "margin", "chosen_reward", "rejected_reward"}
    averaged over pairs; a reward is beta * (policy - ref) and accuracy is the
    share of pairs whose chosen reward beats the rejected one."""
    raise NotImplementedError