roeeaharoni

import torch
from pytorch_pretrained_bert import BertForMaskedLM, BertTokenizer
import random

# Requires pytorch_pretrained_bert: https://github.com/huggingface/pytorch-pretrained-BERT

# returns the probabilities over the vocabulary for the masked words in sent
def get_preds(sent):
    tokenized = bert_tokenizer.tokenize(sent)
    tokenized = ['[CLS]'] + ['[MASK]' if x == 'mask' else x for x in tokenized] + ['[SEP]']