app.py

import gradio as gr
import onnxruntime as rt
from transformers import AutoTokenizer
import torch, json

tokenizer = AutoTokenizer.from_pretrained("distilroberta-base")

with open("genre_types_encoded.json", "r") as fp:
  encode_genre_types = json.load(fp)

genres = list(encode_genre_types.keys())

inf_session = rt.InferenceSession('imdb-classifier-quantized.onnx')
input_name = inf_session.get_inputs()[0].name
output_name = inf_session.get_outputs()[0].name

def imdb_genre(description):
  input_ids = tokenizer(description)['input_ids'][:512]
  logits = inf_session.run([output_name], {input_name: [input_ids]})[0]
  logits = torch.FloatTensor(logits)
  probs = torch.sigmoid(logits)[0]
  return dict(zip(genres, map(float, probs))) 

examples = [
    'Two years of nights have turned Bruce Wayne into a nocturnal animal. But as he continues to find his way as Gothams dark knight Bruce is forced into a game of cat and mouse with his biggest threat so far, a manic killer known as "The Riddler" who is filled with rage and determined to expose the corrupt system whilst picking off all of Gothams key political figures',
    'Infinity has different definitions but one idea the idea of no end. The film is based on this concept as the events are played backwards presenting an endless loop of events which have been lived and experienced by a group of teenagers who represent the majority of those who started the revolution; hoping for change and searching for a better future.',
    'At a fading vacation resort, 11-year-old Sophie treasures rare time together with her loving and idealistic father, Calum (Paul Mescal). As a world of adolescence creeps into view, beyond her eye Calum struggles under the weight of life outside of fatherhood'

    ]


label = gr.outputs.Label(num_top_classes=5)
iface = gr.Interface(fn=imdb_genre, inputs="text", outputs=label, examples = examples)
iface.launch(inline=False)