def postprocess_token_labels(
    labels, input_mask, label_map=None, remove_trailing_word_pieces=False, trailing_token_mask=None
):
    """
    Postprocesses token classification output:
        1) Removes predictions on padded tokens.
        2) If label_map is provided, maps predicted numerical labels
            back to original labels.
        3) If remove_trailing_word_pieces is True and trailing_token_mask
            is provided, remove the predicted labels on trailing word pieces
            generated by WordPiece tokenizer.

    Args:
        labels (list): List of lists of predicted token labels.
        input_mask (list): List of lists. Each sublist contains the attention
            mask of the input token list, 1 for input tokens and 0
            for padded tokens.
        label_map (dict, optional): A dictionary mapping original labels
            (which may be string type) to numerical label ids. If
            provided, it's used to map predicted numerical labels back to
            original labels. Default value is None.
        remove_trailing_word_pieces (bool, optional): Whether to remove
            predicted labels of trailing word pieces generated by WordPiece
            tokenizer. For example, "criticize" is broken into "critic" and
            "##ize". After removing predicted label for "##ize",
            the predicted label for "critic" is assigned to the original word
            "criticize". Default value is False.
        trailing_token_mask (list, optional): list of boolean values, True for
            the first word piece of each original word, False for trailing
            word pieces, e.g. ##ize. If remove_trailing_word_pieces is
            True, this mask is used to remove the predicted labels on
            trailing word pieces, so that each original word in the input
            text has a unique predicted label.
    """
    if label_map:
        reversed_label_map = {v: k for k, v in label_map.items()}
        labels_org = [[reversed_label_map[l_i] for l_i in l] for l in labels]
    else:
        labels_org = labels

    labels_org_no_padding = [
        [label for label, mask in zip(label_list, mask_list) if mask == 1]
        for label_list, mask_list in zip(labels_org, input_mask)
    ]

    if remove_trailing_word_pieces and trailing_token_mask:
        # Remove the padded values in trailing_token_mask first
        token_mask_no_padding = [
            [token for token, padding in zip(t_mask, p_mask) if padding == 1]
            for t_mask, p_mask in zip(trailing_token_mask, input_mask)
        ]

        labels_no_trailing_pieces = [
            [label for label, mask in zip(label_list, mask_list) if mask]
            for label_list, mask_list in zip(labels_org_no_padding, token_mask_no_padding)
        ]
        return labels_no_trailing_pieces
    else:
        return labels_org_no_padding

    Postprocesses token classification output:
        1) Removes predictions on padded tokens.
        2) If label_map is provided, maps predicted numerical labels
            back to original labels.
        3) If remove_trailing_word_pieces is True and trailing_token_mask
            is provided, remove the predicted labels on trailing word pieces
            generated by WordPiece tokenizer.

    Args:
        labels (list): List of lists of predicted token labels.
        input_mask (list): List of lists. Each sublist contains the attention
            mask of the input token list, 1 for input tokens and 0
            for padded tokens.
        label_map (dict, optional): A dictionary mapping original labels
            (which may be string type) to numerical label ids. If
            provided, it's used to map predicted numerical labels back to
            original labels. Default value is None.
        remove_trailing_word_pieces (bool, optional): Whether to remove
            predicted labels of trailing word pieces generated by WordPiece
            tokenizer. For example, "criticize" is broken into "critic" and
            "##ize". After removing predicted label for "##ize",
            the predicted label for "critic" is assigned to the original word
            "criticize". Default value is False.
        trailing_token_mask (list, optional): list of boolean values, True for
            the first word piece of each original word, False for trailing
            word pieces, e.g. ##ize. If remove_trailing_word_pieces is
            True, this mask is used to remove the predicted labels on
            trailing word pieces, so that each original word in the input
            text has a unique predicted label.
    
utils_nlp/models/bert/token_classification.py