EuroPython2022/Mask-Language-Modeling-Using-Pytorch

1import gradio as gr2import torch3import transformers4from transformers import BertTokenizer, BertForMaskedLM5 6device = torch.device('cpu')7 8NUM_CLASSES=59 10model=BertForMaskedLM.from_pretrained("./")11tokenizer=BertTokenizer.from_pretrained("./")12 13 14def predict(text=None) -> dict:  15    model.eval()16    inputs = tokenizer(str(text), return_tensors="pt")17    input_ids = inputs["input_ids"].to(device)18    attention_mask = inputs["attention_mask"].to(device)19    model.to(device)20    token_logits = model(input_ids, attention_mask=attention_mask).logits21    mask_token_index = torch.where(inputs["input_ids"] == tokenizer.mask_token_id)[1]22    mask_token_logits = token_logits[0, mask_token_index, :]23    top_5_tokens = torch.topk(mask_token_logits, NUM_CLASSES, dim=1).indices[0].tolist()24    score = torch.nn.functional.softmax(mask_token_logits)[0]25    top_5_score = torch.topk(score, NUM_CLASSES).values.tolist()26    return {tokenizer.decode([tok]): float(score) for tok, score in zip(top_5_tokens, top_5_score)}27    28gr.Interface(fn=predict, 29             inputs=gr.inputs.Textbox(lines=2, placeholder="Your Text… "),30             title="Mask Language Modeling",31             outputs=gr.outputs.Label(num_top_classes=NUM_CLASSES),32             description="Masked language modeling is the task of masking some of the words in a sentence and predicting which words should replace those masks",33             examples=['A Good Man Is Hard to Find [MASK].', 'Some stories have a [MASK] kind of message called a moral.'],34             interpretation='default').launch()