entertainment genres app

app.py

@@ -1,9 +1,65 @@
 import gradio as gr
+import json
+import torch
+from transformers import AutoTokenizer
+import onnxruntime as rt
+import platform
 
 
-def greet(name):
-    return "Hello " + name + "!!"
+if platform.system() == "Windows":
+    import pathlib
+    temp = pathlib.PosixPath
+    pathlib.PosixPath = pathlib.WindowsPath
 
+model_path = "entertainment-genre-quantized.onnx"
 
-iface = gr.Interface(fn=greet, inputs="text", outputs="text")
-iface.launch()
+with open("genre_types_encoded.json", "r") as file:
+    categories = json.load(file)
+
+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 get_top_label(cat_dict, idx):
+    for key, value in cat_dict.items():
+        if idx == value:
+            return key
+
+
+def get_top_probs(cat_probs, idx):
+    return cat_probs[idx]
+
+
+def entertainment_genres(description):
+    input_ids = tokenizer(description)['input_ids'][:512]
+    probs = inf_session.run([output_name], {input_name: [input_ids]})[0]
+    top_3_indices = sorted(range(len(probs[0])), key=lambda idx: probs[0][idx], reverse=True)[:3]
+    cat_prob = torch.sigmoid(torch.FloatTensor(probs))[0]
+    print(cat_prob)
+
+    top_labels = []
+    for i in top_3_indices:
+        top_labels.append(get_top_label(categories, i))
+
+    top_probs = []
+    for i in top_3_indices:
+        top_probs.append(get_top_probs(cat_prob, i))
+
+    return dict(zip(top_labels, map(float, top_probs)))
+
+
+example = [
+    ["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"],
+    ["When the menace known as the Joker wreaks havoc and chaos on the people of Gotham, Batman must accept one of "
+     "the greatest psychological and physical tests of his ability to fight injustice."]
+]
+
+
+label = gr.outputs.Label(num_top_classes=3)
+
+iface = gr.Interface(fn=entertainment_genres, inputs="text", outputs=label, examples=example)
+iface.launch(inline=False)

requirements.txt

@@ -1 +1,4 @@
-gradio
+gradio==3.44.0
+torch==2.0.1
+transformers==4.33.1
+onnxruntime==1.15.1