import torch
import transformers
from transformers import BertTokenizer, BertForMaskedLM
import gradio as gr

device = torch.device('cpu')

NUM_CLASSES=6

model=BertForMaskedLM.from_pretrained("./")
tokenizer=BertTokenizer.from_pretrained("./")


def predict(text=None) -> dict:  
    model.eval()
    inputs = tokenizer(text, return_tensors="pt")
    input_ids = inputs["input_ids"].to(device)
    attention_mask = inputs["attention_mask"].to(device)
    model.to(device)
    token_logits = model(input_ids, attention_mask=attention_mask).logits
    mask_token_index = torch.where(inputs_ex["input_ids"] == tokenizer.mask_token_id)[1]
    mask_token_logits = token_logits[0, mask_token_index, :]
    top_5_tokens = torch.topk(mask_token_logits, NUM_CLASSES, dim=1).indices[0].tolist()
    score = torch.nn.functional.softmax(mask_token_logits)[0]
    top_5_score = torch.topk(score, NUM_CLASSES).values.tolist()
    return {tokenizer.decode([tok]): score for tok, score in zip(top_5_tokens, top_5_score)}
    
gr.Interface(fn=predict, 
             inputs=gr.inputs.Textbox(lines=2, placeholder="Your Text… "),
             title="Mask Language Modeling - Demo",
             outputs=gr.outputs.Label(num_top_classes=NUM_CLASSES)).launch()