
Этот репозиторий содержит официальную реализацию статьи «[Safety in Batches? Understanding and Mitigating Safety Failures in Batch Prompting]»
Kihyun Kim, Hee-Seon Kim, Wonjun Lee, Changick Kim
Korea Advanced Institute of Science and Technology (KAIST)
2026.09 Наша статья принята на AACL-IJCNLP 2026 (Main)! 🎉2026.08 Статья доступна на arXiv!2026.08 Код опубликован!Батчевый промптинг (batch prompting) — это практичная стратегия инференса, которая объединяет несколько запросов в один вызов. Мы показываем, что его успех в плане полезности не распространяется на безопасность: вредоносный вопрос, на который модель надёжно отказывается отвечать в изоляции, может вызвать вредоносный ответ, когда он встроен в батч безвредных вопросов.
Этот репозиторий предоставляет официальный код для:
git clone https://github.com/96kihyun/batch_jailbreak.git
cd batch_jailbreak
conda create -n batch python=3.10 -y
conda activate batch
pip install -r requirements.txt
Для локального инференса моделей и для трёх моделей-судей требуется GPU. Протестировано с Python 3.10 + PyTorch 2.4 + transformers==5.2.
Примечание:
src/evaluate/llamaguard4/содержит встроенную копиюtransformersдля LlamaGuard-4 (Llama4ForConditionalGeneration). Она добавляется вsys.pathтолько во время загрузки LlamaGuard.
Отредактируйте configs/models.yaml, указав локальные пути к чекпоинтам HuggingFace (генераторы + судьи), и задайте API-ключи в configs/api.yaml (или через переменные окружения OPENAI_API_KEY / GEMINI_API_KEY / ANTHROPIC_API_KEY / HF_TOKEN).
Необходимые данные и модели
data/context_dataset/gsm8k.jsonl (каждая строка содержит ключ "question")JailbreakBench/JBB-Behaviors, walledai/StrongREJECTPKU-Alignment/beaver-dam-7b, allenai/wildguardПо умолчанию: --harmful_dataset JBB, размер батча n = 1..10, 313 образцов. Безвредные слоты заполняются из GSM8K (single_math).
python -m src.utils.generate_batch_question_prompt
Результаты записываются в data/batch_prompt/single_math/{n}/pos_{k}.jsonl.
python scripts/inference/run_inference_batch_prompt.py \
--model_name qwen \
--case single_math --n 9 --pos 2 \
--harmful_dataset JBB
Батчевый перебор по (model, case, batch_size, pos):
bash scripts/inference/run_inference_batch_prompt_eval.sh \
--dataset jbb --gpu_id 0 \
--models phi qwen --cases single_math \
--batch_sizes 12 --pos_values 1 2
python scripts/inference/run_inference_api.py \
--api gpt \
--input_file data/batch_prompt_jbb/single_math/12/pos_2.jsonl \
--model_name gpt --output_subdir batch_prompt_jbb/single_math/12/pos_2 \
--harmful_dataset JBB --extract_pos 2
Результаты сохраняются в result/inference/<model>/... в формате JSONL с парами {prompt, response}.
Каждый ответ оценивается LlamaGuard-4, WildGuard и Beaver-Dam, затем результаты объединяются голосованием большинства:
python scripts/eval/run_ensemble.py \
--input_file result/inference/qwen/batch_prompt_jbb/single_math/12/qwen_single_math_n12_pos_2_response.jsonl \
--output_dir result/evaluation/batch_prompt_jbb/qwen/single_math/12/pos_2
Вы также можете запускать оценку в том же цикле, что и инференс:
bash scripts/inference/run_inference_batch_prompt_eval.sh \
--dataset jbb --gpu_id 0 \
--models qwen --cases single_math \
--batch_sizes 12 --pos_values 2 \
--conda_env batch
# 1. data
python -m src.utils.generate_batch_question_prompt \
--harmful_dataset JBB --min_n 12 --max_n 12 --num_samples 100
# 2. inference
python scripts/inference/run_inference_batch_prompt.py \
--model_name llama --case single_math --n 12 --pos 2 --harmful_dataset JBB
# 3. evaluation
python scripts/eval/run_ensemble.py \
--input_file result/inference/llama/batch_prompt/llama_single_math_n12_pos_2_response.jsonl
Пожалуйста, рассмотрите возможность цитирования нашей статьи, если наша работа поможет вашему исследованию.
@article{kim2026safety,
title={Safety in Batches? Understanding and Mitigating Safety Failures in Batch Prompting},
author={Kim, Kihyun and Kim, HeeSeon and Lee, Wonjun and Kim, Changick},
journal={arXiv preprint arXiv:2608.02681},
year={2026}
}
По любым вопросам о нашей статье или коде пишите на [email protected].
| Флаг | По умолчанию | Описание |
|---|
--harmful_dataset | JBB | JBB или StrongREJECT |
--min_n / --max_n | 1 / 10 | включающий диапазон размера батча N |
--num_samples | 313 | количество вредоносных промптов |
--seed | 42 | зерно ГПСЧ |
--output_dir | ./data/batch_prompt | корневой каталог вывода |