anime-or-not / app.py
daspartho's picture
typo fix
d9f9642
raw
history blame contribute delete
No virus
986 Bytes
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()