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()