# get the length of each box
tokens_length_list: List[int] = [len(l) for l in input_ids_list]

box_end_token_indices = np.array(list(itertools.accumulate(tokens_length_list)))
box_start_token_indices = box_end_token_indices - np.array(tokens_length_list)