Based on your suggestion, this is what I did.
def group_texts(examples):
        # Concatenate all texts.
        second_sentence_start_pos = [i for i, n in enumerate(examples['input_ids']) if n == 985][1]
        token_type_ids = examples["token_type_ids"].copy()
        labels = examples["input_ids"].copy()

        token_type_ids[:second_sentence_start_pos+1] = [0]*(second_sentence_start_pos+1)
        token_type_ids[second_sentence_start_pos+1:] = [1]*(len(examples['input_ids'])-second_sentence_start_pos)
        labels[:second_sentence_start_pos+1] = [-100]*(second_sentence_start_pos+1)

        examples["token_type_ids"] = token_type_ids
        examples["labels"] = labels
        print(examples)
        return examples

Is this what you meant ? Is this the right way to deal with attention_masks and token_type_ids ?