Build next-token training labels that supervise only assistant content. At source position t, the target is the token at position t + 1. Include that target only when both positions are attended and the target role equals assistant_role. Roles describe individual tokens; the source role does not determine eligibility.
Every other label is -100, including the final position. The Boolean loss mask is true exactly where a target is included. Return a Python dictionary with labels, a PyTorch tensor matching the input token shape and dtype, and loss_mask, a Boolean PyTorch tensor of the same shape. Both outputs stay on the input device.
Input: input_ids = [[10,11,20,21,2]], role_ids = [[-1,0,-1,1,-1]], attention_mask = [[true,true,true,true,true]], assistant_role = 1
Output: {"labels":[[-100,-100,21,-100,-100]],"loss_mask":[[false,false,true,false,false]]}
Explanation: Position two predicts token 21, the only next token marked as attended assistant content.
Input: input_ids = [[1,10,11,12]], role_ids = [[-1,0,0,-1]], attention_mask = [[true,true,true,true]], assistant_role = 1
Output: {"labels":[[-100,-100,-100,-100]],"loss_mask":[[false,false,false,false]]}
Input: input_ids = [[1,2,3]], role_ids = [[-1,1,1]], attention_mask = [[true,true,true]], assistant_role = 1
Output: {"labels":[[2,3,-100]],"loss_mask":[[true,true,false]]}
torch.full_like creates an output with the input shape and dtype.
Adjacent sequence slices align positions with their next-token targets.
torch.where selects values using a Boolean tensor.
Accepts: array
Accepts: array
Accepts: array
Accepts: number
Build next-token training labels that supervise only assistant content. At source position t, the target is the token at position t + 1. Include that target only when both positions are attended and the target role equals assistant_role. Roles describe individual tokens; the source role does not determine eligibility.
Every other label is -100, including the final position. The Boolean loss mask is true exactly where a target is included. Return a Python dictionary with labels, a PyTorch tensor matching the input token shape and dtype, and loss_mask, a Boolean PyTorch tensor of the same shape. Both outputs stay on the input device.
Input: input_ids = [[10,11,20,21,2]], role_ids = [[-1,0,-1,1,-1]], attention_mask = [[true,true,true,true,true]], assistant_role = 1
Output: {"labels":[[-100,-100,21,-100,-100]],"loss_mask":[[false,false,true,false,false]]}
Explanation: Position two predicts token 21, the only next token marked as attended assistant content.
Input: input_ids = [[1,10,11,12]], role_ids = [[-1,0,0,-1]], attention_mask = [[true,true,true,true]], assistant_role = 1
Output: {"labels":[[-100,-100,-100,-100]],"loss_mask":[[false,false,false,false]]}
Input: input_ids = [[1,2,3]], role_ids = [[-1,1,1]], attention_mask = [[true,true,true]], assistant_role = 1
Output: {"labels":[[2,3,-100]],"loss_mask":[[true,true,false]]}
torch.full_like creates an output with the input shape and dtype.
Adjacent sequence slices align positions with their next-token targets.
torch.where selects values using a Boolean tensor.
Accepts: array
Accepts: array
Accepts: array
Accepts: number
