Score multiple-choice continuations by their summed conditional log likelihood. Logits and target IDs are already aligned. For batch item b, choice c and token position t, let zbct be the vocabulary-logit vector, xbct the target ID and mbct the Boolean selection mask. The score is:
sbc=t∑mbctlogsoftmax(zbct)xbct.The selected choice is:
cb∗=cargmaxsbc.Here, softmax is over vocabulary entries and log is the natural logarithm. Masked positions contribute zero; do not divide by continuation length. A choice with no selected tokens has score zero, and exact ties select the lowest choice index.
Return a Python dictionary with choice_scores, a PyTorch tensor of shape (batch, choices) matching the logits dtype and device, and predicted_indices, an int64 PyTorch tensor of shape (batch,) on that device.
Input: logits = [[[[2,0]],[[0,2]]]], token_ids = [[[0],[0]]], continuation_mask = [[[true],[true]]]
Output: {"choice_scores":[[-0.126928,-2.126928]],"predicted_indices":[0]}
Explanation: The first choice assigns more probability to its target token and therefore has the larger summed log likelihood.
Input: logits = [[[[1,0],[1,0]],[[2,0],[0,0]]]], token_ids = [[[0,0],[0,1]]], continuation_mask = [[[true,true],[true,false]]]
Output: {"choice_scores":[[-0.626523,-0.126928]],"predicted_indices":[1]}
Input: logits = [[[[0,0]],[[2,-2]],[[-2,2]]]], token_ids = [[[0],[0],[1]]], continuation_mask = [[[false],[false],[false]]]
Output: {"choice_scores":[[0,0,0]],"predicted_indices":[0]}
torch.log_softmax computes log probabilities over a chosen dimension.
Tensor.gather selects each target token from its vocabulary axis.
torch.argmax returns the first index when maximum values tie.
Accepts: array
Accepts: array
Accepts: array
Score multiple-choice continuations by their summed conditional log likelihood. Logits and target IDs are already aligned. For batch item b, choice c and token position t, let zbct be the vocabulary-logit vector, xbct the target ID and mbct the Boolean selection mask. The score is:
sbc=t∑mbctlogsoftmax(zbct)xbct.The selected choice is:
cb∗=cargmaxsbc.Here, softmax is over vocabulary entries and log is the natural logarithm. Masked positions contribute zero; do not divide by continuation length. A choice with no selected tokens has score zero, and exact ties select the lowest choice index.
Return a Python dictionary with choice_scores, a PyTorch tensor of shape (batch, choices) matching the logits dtype and device, and predicted_indices, an int64 PyTorch tensor of shape (batch,) on that device.
Input: logits = [[[[2,0]],[[0,2]]]], token_ids = [[[0],[0]]], continuation_mask = [[[true],[true]]]
Output: {"choice_scores":[[-0.126928,-2.126928]],"predicted_indices":[0]}
Explanation: The first choice assigns more probability to its target token and therefore has the larger summed log likelihood.
Input: logits = [[[[1,0],[1,0]],[[2,0],[0,0]]]], token_ids = [[[0,0],[0,1]]], continuation_mask = [[[true,true],[true,false]]]
Output: {"choice_scores":[[-0.626523,-0.126928]],"predicted_indices":[1]}
Input: logits = [[[[0,0]],[[2,-2]],[[-2,2]]]], token_ids = [[[0],[0],[1]]], continuation_mask = [[[false],[false],[false]]]
Output: {"choice_scores":[[0,0,0]],"predicted_indices":[0]}
torch.log_softmax computes log probabilities over a chosen dimension.
Tensor.gather selects each target token from its vocabulary axis.
torch.argmax returns the first index when maximum values tie.
Accepts: array
Accepts: array
Accepts: array