emotions / app.py
Lord-Raven
Simplifying.
527cc9f
raw
history blame contribute delete
778 Bytes
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()