app.py

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-large")
tokenizer = AutoTokenizer.from_pretrained("roberta-large")
model = AutoModelForMaskedLM.from_pretrained(artist_name)
model2 = AutoModelForMaskedLM.from_pretrained("BigSalmon/FormalRobertaa")



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.')
    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]
        a_list.append(words)
      with torch.no_grad():
        output = model2(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]
        words2 = [tokenizer.decode(i.item()).strip() for i in idx]
        a_list.append(words2)
        st.text_area(label = 'Infill:', value=a_list)