import random


def create_masked_lm_predictions(*, args, tokens, cls_token_id, sep_token_id, mask_token_id,
                                 non_special_ids):
    """Creates the predictions for the masked LM objective."""
    cand_indexes = [i for i, tok in enumerate(tokens) if tok not in (cls_token_id, sep_token_id)]
    output_tokens = list(tokens)
    random.shuffle(cand_indexes)
    num_to_predict = min(args.max_predictions_per_seq,
                         max(1, int(round(len(tokens) * args.masked_lm_prob))))
    mlm_positions = []
    mlm_labels = []
    covered_indexes = set()
    for index in cand_indexes:
        if len(mlm_positions) >= num_to_predict:
            break
        if index in covered_indexes:
            continue
        covered_indexes.add(index)
        masked_token = None

        # 80% of the time, replace with [MASK]
        if random.random() < 0.8:
            masked_token = mask_token_id
        else:
            # 10% of the time, keep original
            if random.random() < 0.5:
                masked_token = tokens[index]
                # 10% of the time, replace with random word
            else:
                masked_token = random.choice(non_special_ids)

        output_tokens[index] = masked_token
        mlm_positions.append(index)
        mlm_labels.append(tokens[index])
    assert len(mlm_positions) <= num_to_predict
    assert len(mlm_positions) == len(mlm_labels)
    return output_tokens, mlm_positions, mlm_labels