import re
import gradio as gr
from transformers import pipeline


generator = pipeline('text-generation', 
                     model='plasticfruits/gpt2-finetuned-how-to-qa', 
                     tokenizer='plasticfruits/gpt2-finetuned-how-to-qa')

def clean_response(user_prompt, response):
    response = re.sub("(?<=\.)[^.]*$", "", response)  # finish at last sentence dot
    response = (
        response.replace("[WP]", "").replace(user_prompt, "").replace("[RESPONSE]", "").replace("<|startoftext|>", "")
    )
    response = response.lstrip()
    return response


def generate(text, length=350):
    prompt = f"\n<|startoftext|>[WP] {text} \n[RESPONSE]"
    result = generator(prompt, 
                       max_length=length, 
                       num_return_sequences=1,
                       do_sample=True,
                       top_k=50,
                       top_p=0.95)
    clean_text = clean_response(text, result[0]["generated_text"])
    return clean_text


examples = [
    ["How to draw a circle"],
    ["How to create a universe"],
    ["How to make pasta"]
]

title = "How-to Generator"
description = "Ask your 'how-to' question to get the best possible answer available in the universe.<br>For best performance, start your question with 'How to {your question}'"
article = "<p style='text-align: center'><a href='how-to-generator.herokuapp.com/' target='_blank'>Official How-To Page</a></p>"


demo = gr.Interface(
    fn=generate,
    inputs=[gr.inputs.Textbox(lines=5, label="Input Text"), gr.Slider(60, 600, value=300, label="Answer Length")],
    outputs=gr.outputs.Textbox(label="Generated Text"),
    examples=examples,
    title=title,
    description=description,
    article=article
)

demo.launch()