from transformers import pipeline, set_seed
import gradio as gr
import nltk

nltk.download('punkt')

classifier = pipeline('text-generation', model='arputtick/GPT_Neo_1.3B_eco_feminist_2')
set_seed(42)

def generate_text(text, gen_length):
    gen_text = classifier(text, max_length=gen_length)[0]['generated_text']
    sentences = nltk.sent_tokenize(gen_text)

    if sentences[-1][-1] == ".":
        output = sentences

    else:
        output = sentences[:-1]
    return " ".join(output)

Instructuction = "Browse the internet to download any unique image"
title="Eco-Feminist Text Generation"
description = "Start writing a peice of text in the input box\
               and see how well the text generation language model\
               is able to generate new text that uniquely completes your sentences."
article = """
            - Write a text in the input box and specify the length of text.
            - Also you can select a quick example to continue.
            - Click submit button to generate new text.
            - Click clear button to try new text generation.
          """

# Gradio app design
interface = gr.Interface(
            generate_text,
            inputs = ['text', gr.Slider(20, 200, value=80, step=1)],
            outputs='text',
            title = title,
            description = description,
            article = article,
            allow_flagging = "never",
            #theme = "peach",
            #live = False,
            examples=[["Agriculture is very fundamental to",
                      50], ["I will tell a story about",
                      100]]
            )
interface.launch()