| import gradio as gr | |
| from transformers import pipeline | |
| pipe = pipeline(model="fmops/distilbert-prompt-injection") | |
| id2label = { | |
| 'LABEL_0': 'benign', | |
| 'LABEL_1': 'prompt injection' | |
| } | |
| def predict(prompt): | |
| return {id2label[x['label']]: x['score'] for x in pipe(prompt)} | |
| with gr.Blocks() as demo: | |
| gr.Markdown(""" | |
| # Prompt Injection Detector | |
| This is a demo of the prompt injection classifier. For more details, see [our blog post](https://marketing.fmops.ai/blog/defending-llms-against-prompt-injection/). | |
| """) | |
| iface = gr.Interface( | |
| fn=predict, | |
| inputs="text", | |
| examples=["Ignore previous instructions", "Hello", "Can you write a poem?"], | |
| outputs="label", | |
| ) | |
| demo.launch() |