import random


def random_word_with_token_ids(token_ids, tokenizer):
    """
    Masking some random tokens for Language Model task with probabilities as in the original BERT paper.
    :param token_ids: list of int, list of token id.
    :param tokenizer: Tokenizer, object used for tokenization (we need it's vocab here)
    :return: (list of str, list of int), masked tokens and related labels for LM prediction
    """
    output_label = []
    mask_id = tokenizer.convert_tokens_to_ids(['[MASK]'])[0]

    for i, token_id in enumerate(token_ids):
        prob = random.random()
        # mask token with 15% probability
        if prob < 0.15:
            prob /= 0.15

            # 80% randomly change token to mask token
            if prob < 0.8:
                token_ids[i] = mask_id

            # 10% randomly change token to random token
            elif prob < 0.9:
                token_ids[i] = random.choice(list(tokenizer.vocab.items()))[1]

            # -> rest 10% randomly keep current token

            # append current token to output (we will predict these later)
            output_label.append(token_id)
        else:
            # no masking token (will be ignored by loss function later)
            output_label.append(-1)

    return token_ids, output_label