mmchowdhury's picture
Update app.py
14b6c7c verified
raw
history blame
932 Bytes
import gradio as gr
import onnxruntime as rt
from transformers import AutoTokenizer
import torch
import json
tokenizer = AutoTokenizer.from_pretrained("distilroberta-base")
with open("tag_types_encoded(1).json", "r") as fp:
encode_genre_types = json.load(fp)
genres = list(encode_genre_types.keys())
inf_session = rt.InferenceSession('Quote-classifier-quantized (3).onnx')
input_name = inf_session.get_inputs()[0].name
output_name = inf_session.get_outputs()[0].name
def classify_Quote_tag(Quote):
input_ids = tokenizer(Quote)['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)))
iface = gr.Interface(
fn=classify_Quote_tag,
inputs="text",
outputs=gr.Label(num_top_classes=5) # Use gr.Label for the label output
)
iface.launch(inline=False)