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


with open("encoded_categories.json", "r") as file:
    categories_encoding = json.load(file)

categories = list(categories_encoding.keys())

tokenizer = AutoTokenizer.from_pretrained("distilroberta-base")
inf_session = rt.InferenceSession('models/manuscript-matcher-quantized.onnx')
input_name = inf_session.get_inputs()[0].name
output_name = inf_session.get_outputs()[0].name


def classify_journal_category(text):
    input_ids = tokenizer(text)['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(categories, map(float, probs)))


label = gr.outputs.Label(num_top_classes=5)
title = "Manuscript Matcher"
description = "<p align=center>This is a demo to classify research articles based on the abstract.</p>"

iface = gr.Interface(fn=classify_journal_category,
                     inputs="text",
                     outputs=label,
                     title=title,
                     description=description,
)
iface.launch(inline=False)