from transformers import pipeline
import gradio as gr

PIPELINES = {}


def build_pipeline(size):
    global PIPELINES
    if size in PIPELINES:
        return PIPELINES[size]

    PIPELINES[size] = pipeline(
        "text2text-generation", model=f"google/flan-t5-{size}", max_length=256
    )
    return PIPELINES[size]


def greet(input_text, size):
    pipe = build_pipeline(size)
    return pipe(input_text)[0]["generated_text"]


demo = gr.Interface(
    fn=greet,
    inputs=[
        gr.Textbox(lines=2, placeholder="Enter your task text..."),
        gr.Radio(choices=["small", "base", "large", "xl"], value="base"),
    ],
    outputs=[gr.Textbox(lines=2)],
)

if __name__ == "__main__":
    demo.launch()