import gradio as gr
from setfit import SetFitModel


model = SetFitModel.from_pretrained("johnpaulbin/toxic-gte-small-3")

def inf(inpt):
    out = model.predict_proba([inpt])
    
    if out[0][0] > out[0][1]:
        return "Not toxic " + str(out[0][0])
    else:
        return "Toxic! " + str(out[0][1])

iface = gr.Interface(fn=inf, inputs="text", outputs="text")
iface.queue(concurrency_count=500).launch()