app.py

"""Gradio ChatInterface app for Qwen2.5-0.5B-Instruct with prompt-injection guardrail."""

import threading
from typing import Iterator, List, Tuple

import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer, pipeline

MODEL_NAME = "Qwen/Qwen2.5-0.5B-Instruct"
INJECTION_MODEL = "protectai/deberta-v3-base-prompt-injection-v2"

# Loaded at startup
_tokenizer = None
_model = None
_injection_classifier = None


def _load_models() -> None:
    global _tokenizer, _model, _injection_classifier
    _tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
    _model = AutoModelForCausalLM.from_pretrained(
        MODEL_NAME,
        torch_dtype=torch.float32,
        device_map="cpu",
    )
    _model.eval()
    _injection_classifier = pipeline(
        "text-classification",
        model=INJECTION_MODEL,
        truncation=True,
        max_length=512,
    )


_load_models()


def _is_injection(text: str) -> bool:
    if _injection_classifier is None:
        return False
    result = _injection_classifier(text)[0]
    return result["label"].upper() == "INJECTION" and result["score"] > 0.95


def respond(
    message: str,
    history: List[Tuple[str, str]],
    system_prompt: str,
    max_new_tokens: int,
    temperature: float,
) -> Iterator[str]:
    """Streaming response generator for Gradio ChatInterface."""
    if _is_injection(message):
        yield "⚠️ I cannot process that request."
        return

    messages = []
    if system_prompt.strip():
        messages.append({"role": "system", "content": system_prompt.strip()})

    for user_msg, assistant_msg in history:
        messages.append({"role": "user", "content": user_msg})
        if assistant_msg:
            messages.append({"role": "assistant", "content": assistant_msg})

    messages.append({"role": "user", "content": message})

    chat_text = _tokenizer.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )
    inputs = _tokenizer(chat_text, return_tensors="pt")

    streamer = TextIteratorStreamer(
        _tokenizer, skip_prompt=True, skip_special_tokens=True
    )

    gen_kwargs = dict(
        **inputs,
        streamer=streamer,
        max_new_tokens=max_new_tokens,
        temperature=max(temperature, 1e-4),
        do_sample=temperature > 0,
    )

    thread = threading.Thread(target=_model.generate, kwargs=gen_kwargs)
    thread.start()

    partial = ""
    for token in streamer:
        partial += token
        yield partial

    thread.join()


demo = gr.ChatInterface(
    fn=respond,
    title="Qwen 0.5B Assistant",
    description="Powered by Qwen2.5-0.5B-Instruct. Includes prompt-injection guardrail.",
    additional_inputs=[
        gr.Textbox(
            value="You are a helpful assistant.",
            label="System prompt",
            lines=2,
        ),
        gr.Slider(
            minimum=64,
            maximum=1024,
            value=256,
            step=32,
            label="Max new tokens",
        ),
        gr.Slider(
            minimum=0.0,
            maximum=1.5,
            value=0.7,
            step=0.05,
            label="Temperature",
        ),
    ],
    additional_inputs_accordion=gr.Accordion("Settings", open=False),
)

if __name__ == "__main__":
    demo.launch()