import streamlit as st
from aitextgen import aitextgen
from PIL import Image
import pybase64


def main():

    def set_png_as_page_bg(png_file):
        bin_str = get_base64_of_bin_file(png_file)
        page_bg_img = '''
        <style>
        body {
        background-image: url("data:image/png;base64,%s");
        background-size: cover;
        }
        </style>
        ''' % bin_str

        st.markdown(page_bg_img, unsafe_allow_html=True)
        return

    @st.cache(allow_output_mutation=True)
    def get_base64_of_bin_file(bin_file):
        with open(bin_file, 'rb') as f:
            data = f.read()
        return pybase64.b64encode(data).decode()

    ai = aitextgen('trained_model/pytorch_model.bin', config = 'trained_model/config.json')
    set_png_as_page_bg('plato2.jpg')
    st.header("What are you thinking?")
    prompt_text = st.text_input(label = 'Ask Plato')
    st.subheader("What do you think love is?")
    st.subheader("Love is good. Love is fun and great. It is vast. It is vast. It is big. Love and life are realative but the truth is bitter."
    "Sadly no one loves you. Some people don't get it.")
    st.subheader(ai.generate_one(prompt = prompt_text, top_p = 0.9 ))




if __name__ == '__main__':
    main()