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, AutoModelWithLMHead


st.title('KoGPT2 Demo')

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


@st.cache(allow_output_mutation=True)
def load_model():
    tokenizer = AutoTokenizer.from_pretrained("gpt2")
    model = AutoModelWithLMHead.from_pretrained("BigSalmon/MrLincoln10")
    return tokenizer, model

load_model()

g = """
informal english: the vice president is not automatically entitled to the nomination, if they are unpopular.
Translated into the Style of Abraham Lincoln: should they be abhorred, the vice president is not above ( sparring with challengers for the nomination / brawling for the nomination / facing a spirited challenge for the nomination / braving a bout for the nomination ).
Translated into the Style of Abraham Lincoln: the office of vice president is no entitlement to the honor of the nomination,  should their popularity ( fail to measure up / be in question / be lacking ).
Translated into the Style of Abraham Lincoln: without ( popular backing / public adulation / favorable reviews / a passing consensus ), the vice president ( is not above / cannot hope to dodge / should not think themselves above / should not bank on evading / should not harbor any illusions of skirting ) a nomination challenge.

informal english: i am sending this to you to tell you of an idea that i have that i really like.
Translated into the Style of Abraham Lincoln: i write this correspondence to bend your ear to a proposition which arouses my keenest interest.

informal english: railroads keep being overtaken by corporations. it is really bad thing and needs a quality response in a fast way.
Translated into the Style of Abraham Lincoln: railroad after railroad is being absorbed by industrial enterprises, a troubling condition that demands wise and swift action.
Translated into the Style of Abraham Lincoln: every day, another railroad lamentably falls under the sway of monopolistic interests, necessitating swift and sure corrective action.
Translated into the Style of Abraham Lincoln: without fail, one railroad after another succumbs to the tentacles of predatory wealth, inviting expeditious and resolute action.

informal english:"""

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():
        text = tokenizer.encode(prompt)
        myinput, past_key_values = torch.tensor([text]), None
        myinput = myinput
        myinput= myinput.to(device)
        logits, past_key_values = model(myinput, past_key_values = past_key_values, return_dict=False)
        logits = logits[0,-1]
        probabilities = torch.nn.functional.softmax(logits)
        best_logits, best_indices = logits.topk(150)
        best_words = [tokenizer.decode([idx.item()]) for idx in best_indices]
        text.append(best_indices[0].item())
        best_probabilities = probabilities[best_indices].tolist()
        words = []              
        st.write(best_words)