import streamlit as st
from Generate_text import generate, load_tokenizer_and_model

device = "cpu"
tok, model = load_tokenizer_and_model("models/essays", device)

st.markdown("## Генератор сочинений")
st.markdown("Это приложение представляет из себя примитивный генератор сочинений. В качестве основной модели был взят трансформер RuGPT от Сбера https://huggingface.co/ai-forever/rugpt3small_based_on_gpt2. После этого на том же HF был скачан датасет с сочинениями на русском языке, который использовался для finetuning'а (см. ReadMe)")

st.markdown("Для работы с приложением укажите максимальный размер сочинения, который хотите получить. Далее напишите тему. Для генерации сочинения длиной в 200 токенов потребуется подождать 6-7 минут")
st.markdown("Модель rugpt3small_based_on_gpt2 маленькая и датасет сочинений, на котором она дообучалась, тоже маленький, поэтому не стоит рассчитывать, что получившееся сочинение будет хорошо отражать выбранную тему). Но, по крайней мере, текст получается осмысленным и его стиль вполне соответствует стилю написания сочинений.")
st.markdown("Чем меньше макисмальная длина, тем хуже получается итоговый результат. Относительно неплохие сочинения получаются при длине 500+, но по времни такая генерация занимает десятки минут")
with st.columns(3)[1]:
    st.markdown("<img width=200px src='https://ps-static.cdn-tinkoff.ru/static/ai-pushkin/portrait-2021-12-10-12-24-56.png'>", unsafe_allow_html=True)
# ^-- можно показывать пользователю текст, картинки, ограниченное подмножество html - всё как в jupyter


max_len = st.text_area("Максимальная длина")
text = st.text_area("Тема сочинения")
# ^-- показать текстовое поле. В поле text лежит строка, которая находится там в данный момент
thesis = "<s>Тема:" + "«" + text + "»." + "\nСочинение: " 

if max_len != "":
    begin_index = len(thesis)
    max_len = int(max_len)
    generated = generate(model, tok, thesis, device, max_length=max_len, num_beams=10)
    end_index = (generated[0]).find("</s>")
    
    st.markdown(f"Сочинение:\n")
    st.markdown(generated[0][begin_index:end_index])
# выводим результаты модели в текстовое поле, на потеху пользователю