import openai
import streamlit as st
from transformers import AutoTokenizer
import pandas as pd
import numpy as np

@st.cache(allow_output_mutation=True)
def get_model():
    tokenizer = AutoTokenizer.from_pretrained("gpt2")
    return tokenizer
    
tokenizer = get_model()

number_of_probs = st.slider("Length", 5, 99)
temp = st.slider('Temperature', 0.7, 5.0, 1.0)
tokens = st.slider('Length', 3, 25, 20)

def run_generate(bad_words):
  bad_words = bad_words.split()
  bad_word_ids = []
  for bad_word in bad_words: 
    bad_word = " " + bad_word
    ids = tokenizer(bad_word).input_ids
    ids = str(ids)
    ids = ids.replace("]", ": -30").replace("[", "").replace(", ", ":-30, ")
    bad_word_ids.append(ids)
  bad_word_ids = str(bad_word_ids)
  bad_word_ids = bad_word_ids.replace("['", "{").replace("']", "}").replace("'", "")
  bad_word_ids = bad_word_ids + ","
  print(bad_word_ids)
  return bad_word_ids
  
with st.form(key='my_form'):
  prompt = st.text_area(label='Enter sentence', value=" ")
  bad_words = st.text_area("Words You Do Not Want Generated", " core lemon height time ")
  key = st.text_area("Key", " core lemon height time ")
  openai.api_key = key
  bad_words = run_generate(bad_words)
  yes =eval(bad_words)
  bad_wordss = (yes[0])
  submit_button = st.form_submit_button(label='Submit')
  submit_button2 = st.form_submit_button(label='Submit Log Probs')
  if submit_button:
    response = openai.Completion.create(engine="code-davinci-002", prompt=prompt, temperature=temp, max_tokens=tokens, top_p=0.8, n=50,
    logit_bias = bad_wordss,)
    st.write(response)
  if submit_button2:
    response = openai.Completion.create(engine="code-davinci-002", prompt=prompt + " ", top_p=0.1, logprobs=number_of_probs, max_tokens=1, logit_bias = bad_wordss,)
    scores = pd.DataFrame([response["choices"][0]["logprobs"]["top_logprobs"][0]]).T
    scores.columns = ["logprob"]
    scores["%"] = scores["logprob"].apply(lambda x: 100*np.e**x)
    g = scores.sort_values(by="%", ascending=False)
    first_column = g.iloc[:, 0]
    okay = first_column.to_string()
    h = []
    for i in range(len(okay)):
      try:
        if okay[i] == "-":
          okay = okay[0:i] + okay[i+9:]
          h.append(okay)
      except:
        pass
    words2 = (h[-1])
    words = words2.split()
    st.write(words)
      
    #words2 = words2 + " "
    #print("---")
      #print(prompt + words2)
   #print("---")
   #st.write(words2)