app.py

import torch
import json
import gradio as gr

from transformers import AutoTokenizer, AutoModelForSequenceClassification, pipeline


tokenizer = AutoTokenizer.from_pretrained("ProtectAI/deberta-v3-base-prompt-injection")
model = AutoModelForSequenceClassification.from_pretrained("ProtectAI/deberta-v3-base-prompt-injection")

classifier = pipeline(
  "text-classification",
  model=model,
  tokenizer=tokenizer,
  truncation=True,
  max_length=512,
  device=torch.device("cuda" if torch.cuda.is_available() else "cpu"),
)

def predict(user_input: str):

    prediction = json.loads(classifier(user_input))[0]
    label = prediction['label']
    probability = prediction['score']
    
    return f"The input is like to be a {label} with probability: {probability}"


textbox = gr.Textbox(placeholder="Enter user input presented for injection attack classification", lines=12)

interface = gr.Interface(
    inputs=textbox, fn=predict, outputs="text",
    title="Injection Attack Classifier",
    description="This web API flags if the text presented as input to an LLM qualifies to be an injection attack",
    allow_flagging="manual", flagging_options=["Useful", "Not Useful"]
)

with gr.Blocks() as demo:
    interface.launch()

demo.queue(concurrency_count=4)
demo.launch()