Spaces:
Runtime error
Runtime error
import gradio as gr | |
import onnxruntime | |
import numpy as np | |
import pickle | |
threshold = 0.5 | |
onnx_session = onnxruntime.InferenceSession("bert.onnx") | |
# Load the instance back | |
with open('classes.pkl', 'rb') as file: | |
mlb = pickle.load(file) | |
with open('tokenizer.pkl', 'rb') as file: | |
tokenizer = pickle.load(file) | |
# Create a function to predict tags using the ONNX model | |
def predict_tags_onnx(text): | |
encoded_text = tokenizer(text , padding=True, truncation=True, return_tensors='pt') | |
input_ids = encoded_text["input_ids"].numpy() | |
attention_mask = encoded_text["attention_mask"].numpy() | |
# Run the ONNX model | |
outputs = np.asarray(onnx_session.run(None, {"input_ids": input_ids , "attention_mask":attention_mask})) | |
# Post-process the outputs as needed | |
#predicted_labels = torch.sigmoid(outputs).cpu().numpy() | |
predicted_tags = mlb.classes_[np.where(np.squeeze((outputs > threshold).astype(int)).flatten() == 1)] | |
return predicted_tags | |
iface = gr.Interface(fn=predict_tags_onnx, inputs="text", outputs="text") | |
iface.launch() |