import gradio as gr
from setfit import SetFitModel


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

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

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