import time

import streamlit as st
from transformers import pipeline


@st.cache(allow_output_mutation=True)
def get_generator():
    return pipeline("text-generation", model="sberbank-ai/rugpt3small_based_on_gpt2")


g = get_generator()

st.markdown("**Генератор смешных и не очень текстов**")

input_test = st.text_area("Введите начало текста:")
length = st.slider("Отрегeлируйте длину текста:", min_value=50, max_value=200)
button = st.button("Сгенерировать!")

if button:
    res = g(
        input_test,
        do_sample=True,
        max_length=length,
        top_k=50,
        top_p=0.95,
        num_returned_sequences=1,
    )
    st.markdown(res[0]["generated_text"])