import streamlit as st
import torch
import numpy as np
import transformers 
import random

def load_model():
    model_finetuned = transformers.AutoModelWithLMHead.from_pretrained(
        'tinkoff-ai/ruDialoGPT-small',
        output_attentions = False,
        output_hidden_states = False
    )
    model_finetuned.load_state_dict(torch.load('GPT_sonnik_only.pt', map_location=torch.device('cpu')))
    tokenizer = transformers.AutoTokenizer.from_pretrained('tinkoff-ai/ruDialoGPT-small')
    return model_finetuned, tokenizer

def preprocess_text(text_input, tokenizer):
    prompt = tokenizer.encode(text_input, return_tensors='pt')

def predict_sentiment(model, prompt, temp, num_generate):
    result = model.generate(
        input_ids=prompt,
        max_length=150,
        num_beams=5,
        do_sample=True,
        temperature=float(temp),
        top_k=50,
        top_p=0.6,
        no_repeat_ngram_size=3,
        num_return_sequences=num_generate,
        ).cpu().numpy()
    return result

st.title('Text generation with dreambook')

model, tokenizer = load_model()

text_input = st.text_input("Enter some text about movie")
max_len = st.slider('Length of sequence', 0, 500, 250)
temp = st.slider('Temperature', 0, 30, 0)
if st.button('Generate a random number of sequences'):
    num_generate = random.randint(1,5)
    st.write(f'Number of sequences: {num_generate}')
else:
    num_generate = st.text_input("Enter number of sequences")

if text_input:
    prompt = preprocess_text(text_input, tokenizer)
    result = predict_sentiment(model, prompt, max_len, temp, num_generate)
    for i in result:
        st.write(textwrap.fill(tokenizer.decode(i), max_len))
        if st.button('Next'):
            try:
                continue
            except:
                st.write('All sequences are return. Generate new')
        else:
            None