app.py

import gradio as gr
import transformers

from transformers import AutoModelForSequenceClassification, AutoTokenizer
from transformers import pipeline

title = "Description to Pokemon-Type"
description = """
<center>
Given a written description of a character, real or imagined, living or dead - this bot imagines them as a pokemon
and then describes which pokemon type they would be. For example...

"Born and raised in the Austrian Empire, Nikolai Tesla studied engineering and physics in the 1870s without receiving a degree, gaining practical experience in the early 1880s working in telephony and at Continental Edison in the new electric power industry."
Nikolai Tesla would be an electric-type pokemon.

</center>
"""

article = "Inspired by [this article](https://medium.com/analytics-vidhya/predicting-pok%C3%A9mon-type-with-the-pok%C3%A9dex-7038754dc422)."



model = AutoModelForSequenceClassification.from_pretrained('mrcoombes/distilbert-wikipedia-pokemon')
tokenizer = AutoTokenizer.from_pretrained('distilbert-base-uncased')

classifier = pipeline('text-classification', model = model, tokenizer=tokenizer, return_all_scores=True)

clf = lambda x: sorted(classifier(x)[0], key=lambda y: y['score'], reverse=True)
gr_labels = lambda list_of_dicts: {d['label']: d['score'] for d in list_of_dicts}

compose = lambda description: gr_labels(clf(description))

demo = gr.Interface(fn=compose, inputs="text", outputs=gr.outputs.Label(num_top_classes=5), title=title, description=description, article=article)

demo.launch()