import gradio as gr

title = "Reformer"

description = "Gradio Demo for Reformer'. 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://arxiv.org/abs/2001.04451' target='_blank'>Reformer: The Efficient Transformer</a></p>"

examples = [
    ['A few months later']
]

reformerc = gr.Interface.load("huggingface/google/reformer-crime-and-punishment")

reformere = gr.Interface.load("huggingface/google/reformer-enwik8")

def inference(text,model):
    if model == "reformer-crime-and-punishment":
        outtext = reformerc(text)
    else:
        outtext = reformere(text)
    return outtext   

gr.Interface(
    inference, 
    [gr.inputs.Textbox(label="Input"),gr.inputs.Dropdown(choices=["reformer-crime-and-punishment","reformer-enwik8"], type="value", default="reformer-enwik8", label="model")], 
    gr.outputs.Textbox(label="Output"),
    examples=examples,
    article=article,
    title=title,
    description=description).launch(enable_queue=True, cache_examples=True)