import streamlit as st
from transformers import GPT2LMHeadModel, GPT2Tokenizer, TextDataset, DataCollatorForLanguageModeling, Trainer, TrainingArguments
import torch
st.header('Генератор плохих подкатов')

tokenizer = GPT2Tokenizer.from_pretrained('SeducerWeight/repo_name')
model = GPT2LMHeadModel.from_pretrained('SeducerWeight/repo_name', from_tf=False)

text = st.text_input("Кринжовое начало")
max_length = st.slider('Выберите длину', 30, 150, 40, 5)
temp = st.slider('Выберите температуру', 1.5, 4.0, 2.0, 0.25)
top_p = st.slider('Выберите top_p', 1.5, 4.0, 2.0, 0.25)
if text and max_length and temp and top_p:
    input_ids = tokenizer.encode(text, return_tensors="pt")

    model.eval()
    with torch.no_grad():
        out = model.generate(input_ids, 
                            do_sample=True,
                            num_beams=3,
                            temperature=temp,
                            top_p=top_p,
                            max_length=max_length,
                            )

    generated_text = list(map(tokenizer.decode, out))
    count = st.number_input('Введите количество подкатов', 1, 5, 1)
    st.write(generated_text[0].split('\n')[count-1])