import gradio as gr
from transformers import AutoModelForSequenceClassification, AutoTokenizer, TextClassificationPipeline

tokenizer = AutoTokenizer.from_pretrained("daspartho/anime-or-not")
model = AutoModelForSequenceClassification.from_pretrained("daspartho/anime-or-not") # i've uploaded the model on HuggingFace :)

label_map = {
    'LABEL_0':"That's not Anime enough!",
    'LABEL_1':"That's an Anime!",
}

pipe = TextClassificationPipeline(model=model, tokenizer=tokenizer)

def classify_text(plot):
    return label_map[pipe(plot)[0]['label']]

iface = gr.Interface(
    description = "Write a plot, and the model will determine whether it's anime enough",
    article = "<p style='text-align: center'><a href='https://github.com/daspartho/anime-or-not' target='_blank'>Github</a></p>",
    fn=classify_text, 
    inputs=gr.inputs.Textbox(label="Type the plot here",), 
    outputs=gr.outputs.Textbox(label='What the model thinks'),
    enable_queue=True,
    )
iface.launch()