create_assistant_response_mask failed?

[xiangliustc]

   I want to know why u set pos = im_start_positions[-2] and pos = im_start_positions[-4] in my opinion ,the second logic is right .Especially when u don't use tools or have different data formats.
    for i in range(inputs.shape[0]):
        sequence = inputs[i]
        

        # Find all im_start positions
        im_start_positions = (sequence == im_start).nonzero().flatten()
        if True:
            pos = im_start_positions[-2]
            if pos + 1 < len(sequence) and sequence[pos + 1] == assistant:
                next_end = sequence[pos:].eq(im_end).nonzero()
                if len(next_end) > 0:
                    end_pos = pos + next_end[0].item()
                    # Mark the entire response (including the im_start and im_end tokens)
                    mask[i, pos:end_pos + 1] = True
            pos = im_start_positions[-4]
            if pos + 1 < len(sequence) and sequence[pos + 1] == assistant:
                next_end = sequence[pos:].eq(im_end).nonzero()
                if len(next_end) > 0:
                    end_pos = pos + next_end[0].item()
                    # Mark the entire response (including the im_start and im_end tokens)
                    mask[i, pos:end_pos + 1] = True
        else:
            for pos in im_start_positions:
                # Check if the token after im_start is "assistant"
                if pos + 1 < len(sequence) and sequence[pos + 1] == assistant:
                # Find the next im_end
                    next_end = sequence[pos:].eq(im_end).nonzero()
                    if len(next_end) > 0:
                        end_pos = pos + next_end[0].item()
                        # Mark the entire response (including the im_start and im_end tokens)
                        mask[i, pos:end_pos + 1] = True
                        
                    # Debug print

        # if if_use_weighted:
        #     final_answer_positions = (sequence == answer).nonzero().flatten()
        #     if len(final_answer_positions) == 3:
        #         pos = final_answer_positions[-1]
        #         weighted_mask[i, pos+2] = True

    
    return mask

[cjakfskvnad]

This part sets text token masks corresponding to only the final answer and the correct visual operation to true.
The im_start_positions[-2] sets the text token masks corresponding to the final answer. For example, agent answers :"So the final answer is A"
The im_start_positions[-4] sets the text token masks corresponding to the correct visual operation (the last one) to true. Since we have self correction trajectories: the first visual operations of these trajectories are designed to be inaccurate, we set false masks to these visual operations.