import streamlit as st
import numpy as np
import pandas as pd
import os
import torch
import torch.nn as nn
from transformers.activations import get_activation
from transformers import AutoTokenizer, AutoModelForCausalLM
from transformers import AutoTokenizer, AutoModel
from transformers import GPTNeoXForCausalLM, GPTNeoXTokenizerFast
import math
import numpy as np


st.title('GPT2:')

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

@st.cache(allow_output_mutation=True)
def get_model():
    tokenizer = AutoTokenizer.from_pretrained("BigSalmon/InformalToFormalLincoln85Paraphrase")
    model = AutoModelForCausalLM.from_pretrained("BigSalmon/InformalToFormalLincoln85Paraphrase")
    return model, tokenizer
    
model, tokenizer = get_model()

g = """***

original: sports teams are profitable for owners. [MASK], their valuations experience a dramatic uptick. 
infill: sports teams are profitable for owners. ( accumulating vast sums / stockpiling treasure / realizing benefits / cashing in / registering robust financials / scoring on balance sheets ), their valuations experience a dramatic uptick. 

***

original:"""


def score(tokens_tensor):
    loss=model(tokens_tensor, labels=tokens_tensor)[0]
    return np.exp(loss.cpu().detach().numpy())

def prefix_format(sentence):
  words = sentence.split()
  if "[MASK]" in sentence:
    words2 = words.index("[MASK]")
    #print(words2)
    output = ("<|SUF|> " + ' '.join(words[words2+1:]) + " <|PRE|> " + ' '.join(words[:words2]) + " <|MID|>")
    st.write(output)
  else:
    st.write("Add [MASK] to sentence")
    
with st.form(key='my_form'):
    prompt = st.text_area(label='Enter sentence', value=g)
    submit_button = st.form_submit_button(label='Submit')
    if submit_button:
      with torch.no_grad():
        tokens_tensor = tokenizer.encode(prompt, add_special_tokens=False, return_tensors="pt")
        perplex = score(tokens_tensor)       
        st.write(perplex)