entertainment category app ready

app.py

@@ -1,9 +1,38 @@
 import gradio as gr
+import numpy as np
+import json
+import torch
+from transformers import AutoTokenizer
+import onnxruntime as rt
 
 
-def greet(name):
-    return "Hello " + name + "!!"
+model_path = "entertainment-category-quantized.onnx"
+with open("category.json", "r") as file:
+    categories = json.load(file)["categories"]
 
 
-iface = gr.Interface(fn=greet, inputs="text", outputs="text")
-iface.launch()
+inf_session = rt.InferenceSession(model_path)
+input_name = inf_session.get_inputs()[0].name
+output_name = inf_session.get_outputs()[0].name
+
+tokenizer = AutoTokenizer.from_pretrained("distilroberta-base")
+
+
+def entertainment_category(description):
+    input_ids = tokenizer(description)["input_ids"][:512]
+    probs = inf_session.run([output_name], {input_name: [input_ids]})[0]
+
+    mask = np.where(probs[0] == probs.max())[0][0]
+    cat = categories[mask]
+    cat_prob = torch.sigmoid(torch.FloatTensor(probs))[0]
+
+    return dict(zip(categories, map(float, cat_prob)))
+
+
+with open("example.json", "r") as file:
+    examples = json.load(file)["examples"]
+
+
+label = gr.components.Label(num_top_classes=5)
+iface = gr.Interface(fn=entertainment_category, inputs="text", outputs=label, examples=examples)
+iface.launch(inline=False)

category.json

@@ -0,0 +1 @@
+{"categories": ["anime", "book", "game", "movie", "music", "tv show"]}

example.json

@@ -0,0 +1,5 @@
+{"examples": [
+  ["March Of Soldiers is a real time strategy single player , It is a military game based on the player's skill and the strength of his financial economy"],
+  ["Don Vito Corleone, head of a mafia family, decides to hand over his empire to his youngest son Michael. However, his decision unintentionally puts the lives of his loved ones in grave danger."],
+  ["When I wake up, I look into the mirror, I can see a clearer, vision, I should start living today, Cause today is gonna be the day, is gonna be the day, Cause today is gonna be the day, is gonna be the day"]
+]}