import gradio as gr

from transformers.tools.text_classification import TextClassificationTool

classifier = TextClassificationTool()
labels = ["Star Wars", "Star Trek", "MCU", "LOTR", "Babylon 5", "DC Comics"]

iface = gr.Interface(
    fn=lambda s: classifier(s, labels=labels),
    inputs="text",
    outputs="text")
with gr.Blocks() as app:
    story=gr.TextArea(
        label='Story',
        placeholder='The space pirate drew his trusty blaster. This might be his last night in Mos Eisley')
    output=gr.Textbox(label='Fandom')
    note=gr.Markdown(f"I know {', '.join(labels)}")
    button = gr.Button('Classify')
    button.click(
        fn=lambda s: classifier(s, labels=labels),
        inputs=story,
        outputs=output,
        api_name="fandom")
                 
app.launch()