import gradio as gr

from happytransformer import HappyGeneration

happy_gen = HappyGeneration("GPT2", "DarwinAnim8or/GPT-NoSleep-1.5b")

from happytransformer import GENSettings

def generate(text, length=100, penalty=1, temperature=0.8):
    args_top_k = GENSettings(no_repeat_ngram_size=penalty, do_sample=False, top_k=80, top_p=1, temperature=temperature, max_length=length, early_stopping=False)
    print("Prompt: " + text)
    
    inputText = "[WP] " + text + " [RESPONSE] "
    
    result = happy_gen.generate_text(inputText, args=args_top_k)
    generated_text = result.text #returns generated text only

    generated_text = generated_text.replace('.', '.\n')

    return generated_text

examples = [
    ["We don't go to the forest anymore."],
    ["There used to be a smile on my face. "],
    #["Now I know why they're called 'Mans best friend'"],
    #["The terror in the night"],
    #["I still see them."],
    #["The curse of the plague doctor"]
]

demo = gr.Interface(
    fn=generate,
    inputs=[
        gr.inputs.Textbox(lines=5, label="Input Text"),
        gr.inputs.Slider(25, 500, label='Length', default=250, step=25),
        gr.inputs.Slider(1, 10, label='no repeat ngram size', default=1, step=1),
        gr.inputs.Slider(0.0, 1.0, label='Temperature - control randomness', default=0.6, step=0.1)
    ],
    outputs=gr.outputs.Textbox(label="Generated Text"),
    examples=examples,
    title="NoSleep Horror Story Generator",
    description="Using the 1.5b size model. You may need to run it a few times in order to get something good! TIP: if you want a writing prompt, use this model in combination: https://huggingface.co/spaces/DarwinAnim8or/NoSleepWritingPromptGenerator"
)

demo.launch()