Rimi98/AppClassifier

1import gradio as gr2import onnxruntime 3from transformers import AutoTokenizer4import torch, json5 6token = AutoTokenizer.from_pretrained('distilroberta-base')7 8with open("data.json", "r") as fp:9  types = json.load(fp)10 11types = list(types)12 13 14inf_session = onnxruntime.InferenceSession('classifier-quantized.onnx')15input_name = inf_session.get_inputs()[0].name16output_name = inf_session.get_outputs()[0].name17 18def classify(review):19    input_ids = token(review)['input_ids'][:512]20    logits = inf_session.run([output_name], {input_name: [input_ids]})[0]21    logits = torch.FloatTensor(logits)22    probs = torch.sigmoid(logits)[0]23    return dict(zip(types, map(float, probs))) 24 25 26label = gr.outputs.Label(num_top_classes=5)27iface = gr.Interface(fn=classify, inputs="text", outputs=label)28iface.launch(inline=False)