import streamlit as st
import torch
from transformers import GPT2LMHeadModel, GPT2Tokenizer




model_name_or_path = "sberbank-ai/rugpt3small_based_on_gpt2"
tokenizer = GPT2Tokenizer.from_pretrained(model_name_or_path)
model =  GPT2LMHeadModel.from_pretrained(
    model_name_or_path,
    output_attentions = False,
    output_hidden_states = False,
)


# Загрузка сохраненных весов
model_weights_path = "hunter_pelevin.pt"
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.load_state_dict(torch.load(model_weights_path, map_location=device))
model.eval()
def generate_text(user_input, model=model, tokenizer=tokenizer):
    input_ids = tokenizer.encode(user_input, return_tensors="pt")
    with torch.no_grad():
        out = model.generate(
            input_ids,
            max_length=slider1,
            num_beams=10,
            do_sample=True,
            temperature=slider3,
            top_k=500,
            top_p=0.8,
            no_repeat_ngram_size=3,
            num_return_sequences=slider2,
        )
    generated_text = list(map(tokenizer.decode, out))[0]
    return generated_text


st.title("Простое веб-приложение на Streamlit")
# Получаем ввод от пользователя
user_input = st.text_area("Введите текст:")
slider1 = st.slider("Выберите длинну текста:", min_value=10, max_value=100, value=50)
slider2 = st.slider("Выберите количество генераций", min_value=1, max_value=5, value=2)
slider3 = st.slider("Выберите степень безумия:", min_value=0.1, max_value=3.0, value=1.2, step=0.1)
if user_input:
        gen_text = generate_text(user_input)
        st.write(gen_text)