import streamlit as st
from pages.models.gptmodel import model
from pages.models.gptmodel import tokenizer
import torch
import transformers
from transformers import GPT2LMHeadModel, GPT2Tokenizer

text = st.text_input('Введите сюда вопрос или предложение для генерации текста')


temp = st.slider('Выберите значения температуры',min_value=1.0, max_value=3.0, step=0.1 )

length = st.slider('Выберите длину последовательности',min_value=1, max_value=200, step=10 )

beams = st.slider('Выберите число генераций',min_value=1, max_value=10, step=1 )


p = st.slider('Выберите значения top_p',min_value=0.5, max_value=0.9, step=0.1 )

if text:
    input_ids = tokenizer.encode(text, return_tensors="pt").to('cpu')
    model.eval()
    with torch.no_grad():
        out = model.generate(input_ids,
                            do_sample=True,
                            num_beams=beams,
                            temperature=temp,
                            top_p=p,
                            max_length=length,
                            )

        st.write(list(map(tokenizer.decode, out))[0])