import gradio as gr

from transformers import pipeline

pipe = pipeline(model="fmops/distilbert-prompt-injection")
id2label = {
    'LABEL_0': 'benign',
    'LABEL_1': 'prompt injection'
}

def greet(prompt):
    return {id2label[x['label']]: x['score'] for x in pipe(prompt)}

iface = gr.Interface(fn=greet, inputs="text", outputs="label")
iface.launch()