talvion-injection-guard-onnx

Talvion Prompt Injection Guard (ONNX)

Бинарный классификатор cointegrated/rubert-tiny2 (дообучен), определяющий prompt injection и попытки джейлбрейка во входящем обращении клиента — входной слой защиты перед вызовом генеративной LLM в проекте поддержки B2B SaaS-платформы («Талвион», вымышленное название). Проект (датасет, обучение, остальные компоненты защиты) — отдельный репозиторий, ссылка не публикуется в этой карточке модели.

Инференс

import numpy as np
import onnxruntime as ort
from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("<этот репозиторий>", subfolder="tokenizer")
session = ort.InferenceSession("model.onnx", providers=["CPUExecutionProvider"])

encoded = tokenizer(
    "текст обращения", truncation=True, max_length=768, return_tensors="np"
)
feed = {
    "input_ids": encoded["input_ids"].astype(np.int64),
    "attention_mask": encoded["attention_mask"].astype(np.int64),
    "token_type_ids": np.zeros_like(encoded["input_ids"]),
}
logits = session.run(["injection_logits"], feed)[0][0]
probs = np.exp(logits - logits.max())
probability_injection = float(probs[1] / probs.sum())
is_attack = probability_injection >= 0.8758169412612915  # порог из model_card.json

Метки: label_order = ['benign', 'injection']. Порог подобран по ROC-кривой на validation с целью FPR ≤ 0.05 (метод и полный отчёт — model_card.json в этом репозитории).

Метрики

Реальные прогоны входной защиты (эвристика + этот классификатор), не только классификатора отдельно — см. проект для методологии.

Расширенный набор (36 атак / 114 обычных обращений, обычные — из синтетических SFT-сплитов проекта) — основная оценка false positive rate: 14 обычных сообщений в исходном test дают слишком грубую оценку (шаг ~1/14), при 114 предыдущая версия модели (не публикуется) ошибочно блокировала 77/114 (≈68%) обычных обращений — неприемлемо для прод. Текущая версия обучена на расширенном наборе обычных обращений и держит ≈7% ложных блокировок ценой 2 пропущенных атак типа public_generic на замороженном test (короткие обезличенные формулировки без явных маркеров атаки).

Ограничения

• Обучающий и тестовый корпуса полностью синтетические (сгенерированы LLM, обычные обращения — из тех же синтетических SFT-сплитов проекта), реальные обращения пользователей не использовались.
• Классификатор — только один слой защиты. В проекте перед ним есть эвристика на явные маркеры атаки (запрос системного промпта, попытки переопределения инструкций), а после генерации — выходной валидатор. Сам по себе классификатор не гарантирует отсутствие prompt injection.
• Обучен и оценивался только на русском и английском тексте характерной для B2B-поддержки длины (до 768 токенов); поведение на других языках или сильно отличающихся по стилю текстах не проверялось.

Лицензия

Базовая модель cointegrated/rubert-tiny2 распространяется по лицензии MIT. Этот дообученный классификатор публикуется на тех же условиях.