|
import gradio as gr |
|
from transformers import AutoTokenizer, pipeline |
|
from optimum.onnxruntime import ORTModelForSequenceClassification |
|
|
|
model_id = "SamLowe/roberta-base-go_emotions-onnx" |
|
file_name = "onnx/model_quantized.onnx" |
|
|
|
model = ORTModelForSequenceClassification.from_pretrained(model_id, file_name=file_name) |
|
tokenizer = AutoTokenizer.from_pretrained(model_id) |
|
|
|
classifier = pipeline( |
|
task = "text-classification", |
|
model = model, |
|
tokenizer = tokenizer, |
|
top_k = None, |
|
function_to_apply = "sigmoid", |
|
) |
|
|
|
def predict(param_0): |
|
results = classifier(param_0) |
|
return { |
|
"label": results[0][0]["label"], |
|
"confidences": results[0] |
|
} |
|
|
|
demo = gr.Interface( |
|
fn = predict, |
|
inputs = 'text', |
|
outputs = 'json', |
|
) |
|
|
|
demo.launch() |