import streamlit as st
import transformers
from transformers import pipeline, set_seed


def infer(sent, max_length, num_return_sequences):
    generator = pipeline('text-generation', model='gpt2')
    
    return generator(sent, max_length=max_length, num_return_sequences=num_return_sequences)


st.title("Write with Transformers 🦄")
st.write("The almighty king of text generation, GPT-2 comes in four available sizes, only three of which have been publicly made available. Feared for its fake news generation capabilities, it currently stands as the most syntactically coherent model. A direct successor to the original GPT, it reinforces the already established pre-training/fine-tuning killer duo. From the paper: Language Models are Unsupervised Multitask Learners by Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei and Ilya Sutskever.")

default_value = "See how a modern neural network auto-completes your text 🤗 This site, built by the Hugging Face team, lets you write a whole document directly from your browser, and you can trigger the Transformer anywhere using the Tab key. Its like having a smart machine that completes your thoughts 😀 Get started by typing a custom snippet, check out the repository, or try one of the examples. Have fun!"
sent = st.text_area("Text", default_value, height = 275)
max_length = st.sidebar.slider("Max Length", min_value = 10, max_value=30)
num_return_sequences = st.sidebar.number_input('Number of Sequences to be Generated', min_value=1, max_value=5, value=1, step=1)

outputs = infer(sent, max_length = max_length, num_return_sequences=num_return_sequences)

#for output in outputs[0]:
st.write(outputs[0]["generated_text"])