import gradio as gr
from gradio import mix

title = "GPT2"
description = "Gradio Demo for OpenAI GPT2. To use it, simply add your text, or click one of the examples to load them. Read more at the links below."

article = "<p style='text-align: center'><a href='https://d4mucfpksywv.cloudfront.net/better-language-models/language_models_are_unsupervised_multitask_learners.pdf' target='_blank'>Language Models are Unsupervised Multitask Learners</a></p>"

examples = [
    ['Paris is the capital of']
]

io1 = gr.Interface.load("huggingface/gpt2")

io2 = gr.Interface.load("huggingface/gpt2-large")

io3 = gr.Interface.load("huggingface/gpt2-medium")

io4 = gr.Interface.load("huggingface/gpt2-xl")

def inference(text, model):
    if model == "gpt2-large":
        text = io2(inputs=gr.inputs.Textbox(lines=5, label="Input Text"))
    elif model == "gpt2-medium":
        text = io3(inputs=gr.inputs.Textbox(lines=5, label="Input Text"))
    elif model == "gpt2-xl":
        text = io4(inputs=gr.inputs.Textbox(lines=5, label="Input Text"))
    else:
        text = io1(inputs=gr.inputs.Textbox(lines=5, label="Input Text"))
    return text   
    
     

gr.Interface(
    inference, 
    [gr.inputs.Textbox(label="Input"),gr.inputs.Dropdown(choices=["gpt2-large","gpt2-medium","gpt2-xl"], type="value", default="gpt2-medium", label="model")
], 
    gr.outputs.Textbox(label="Output"),
    examples=examples).launch(debug=True)