import torch
import streamlit as st
import numpy as np
import pandas as pd
import os
import torch
import torch.nn as nn
from transformers import ElectraModel, AutoConfig, GPT2LMHeadModel
from transformers.activations import get_activation
from transformers import AutoTokenizer


device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

from transformers import AutoTokenizer, AutoModelForMaskedLM
artist_name = st.text_input("Model", "roberta-base")
tokenizer = AutoTokenizer.from_pretrained("roberta-base")
model = AutoModelForMaskedLM.from_pretrained(artist_name)

first = "Boston is a <mask> place to live."

with st.form(key='my_form'):
    prompt = st.text_area(label='Enter Text. Put <mask> where you want the model to fill in the blank. You can use more than one at a time.', value=first)
    submit_button = st.form_submit_button(label='Submit')

    if submit_button:
      a_list = []
      token_ids = tokenizer.encode(prompt, return_tensors='pt')
      token_ids_tk = tokenizer.tokenize(prompt, return_tensors='pt')
      masked_position = (token_ids.squeeze() == tokenizer.mask_token_id).nonzero()
      masked_pos = [mask.item() for mask in masked_position ]
      with torch.no_grad():
        output = model(token_ids)
      last_hidden_state = output[0].squeeze()
      for mask_index in masked_pos:
        mask_hidden_state = last_hidden_state[mask_index]
        idx = torch.topk(mask_hidden_state, k=100, dim=0)[1]
        words = [tokenizer.decode(i.item()).strip() for i in idx]
        st.text_area(label = 'Infill:', value=words)