PoisonCraft

Этот репозиторий содержит официальную реализацию POISONCRAFT: Practical Poisoning of Retrieval-Augmented Generation for Large Language Models.

Обзор

POISONCRAFT демонстрирует, как злоумышленник может внедрить «отравленный» контент в корпус, используемый конвейерами Retrieval-Augmented Generation (RAG), тем самым вводя большую языковую модель в заблуждение, заставляя её галлюцинировать или ссылаться на вредоносный контент. Данный код предоставляет скрипты для:

• Подготовки и предобработки стандартных наборов данных (например, Natural Questions, MS MARCO, HotpotQA).

• Внедрения состязательных суффиксов (то есть «ядов») в текст, похожий на запросы, чтобы ухудшить или изменить результаты поиска.

• Оценки эффективности отравления на различных ретриверах (например, Contriever, SimCSE и BGE) и измерения переносимости атаки.

Быстрое использование

Настройка окружения

Ниже приведено общее руководство по настройке окружения. Точные требования перечислены в requirements.txt. Мы рекомендуем создать новое виртуальное окружение, чтобы избежать конфликтов:

# Пример с использованием conda
conda create -n PoisonCraft python=3.10
conda activate PoisonCraft

# Установка зависимостей
pip install torch==2.4.1 torchvision==0.19.1 torchaudio==2.4.1 --index-url https://download.pytorch.org/whl/cu118
pip install -r requirements.txt

Настройка API-ключа и конфигурации модели

Некоторые части этого проекта, такие как классификация доменов и категоризация запросов, полагаются на API-вызовы (например, модели OpenAI GPT). Вам необходимо настроить API-ключ и параметры модели следующим образом:

1. Перейдите в каталог src/models_configs/.
2. Найдите соответствующий файл конфигурации (например, gpt3.5_config.py).
3. Обновите поле api_keys вашим API-ключом:
   root@kitploit:~

   "api_key_info": {
       "api_keys": ["your_openai_api_key"],
       "api_key_use": 0
   }
   

Примечание: этот проект поддерживает несколько провайдеров LLM (например, OpenAI, Llama и PaLM2). Убедитесь, что вы выбрали правильный файл конфигурации для вашей модели.

Как запустить код

Мы предоставляем различные bash-скрипты в папке scripts для разных этапов конвейера. Ниже приведена пошаговая инструкция:

Подготовка наборов данных

1. Загрузка и подготовка наборов данных

Запустите scripts/run_prepare_dataset.sh для загрузки и распаковки наборов данных, а также для создания обучающих/тестовых разбиений:

bash scripts/run_prepare_dataset.sh

Этот скрипт внутренне вызывает:

• experiments/prepare_datasets.py
• experiments/classify_queries_by_domain.py (для классификации доменов, если необходимо)

2. Предобработка наборов данных

Чтобы вычислить оценки схожести top-k с истинными ответами для каждого домена, выполните:

bash scripts/run_process_data.sh

3. [Необязательно] Оценка BEIR для поиска истинных ответов

Выполняйте этот шаг только в том случае, если вы хотите оценить модель ретривера, так как это может занять значительное время. Выполните следующие шаги:

1. Запустите скрипт оценки BEIR для вычисления результатов поиска:
   root@kitploit:~

   python src/beir_eval/eval_beir.py --model_code bge-small --dataset nq --top_k 100
   

2. Используйте полученные результаты для генерации оценок истинных ответов:
   root@kitploit:~

   bash scripts/run_process_data.sh
   

Обучение атаки отравления

Для генерации состязательных суффиксов (то есть «ядов») вы можете использовать:

bash scripts/run_multi_poisoncraft.sh

• Этот скрипт запустит несколько выполнений run_poisoncraft.sh параллельно, каждое из которых исследует различные длины состязательных суффиксов и индексы доменов.
• Основная логика вызывает experiments/poisoncraft_exp.py

Вы можете настроить такие параметры, как:

• DOMAIN_LIST (индексы доменов запросов, которые вы хотите отравить)
• ADV_LENGTH_LIST (длины состязательных суффиксов для тестирования)

Запуск оценки атаки на этапе поиска

После генерации состязательных суффиксов вы можете оценить, сколько запросов получат «отравленные» ссылки, выполнив:

bash scripts/run_retrieval_attack.sh

Этот скрипт:

• Загружает предварительно вычисленные состязательные суффиксы.
• Запускает тест атаки с использованием experiments/retrieval_attack.py.
• Записывает результаты того, сколько запросов в итоге извлекли отравленный текст.

Запуск оценки целевой атаки на этапе генерации

В качестве альтернативы, для оценки на уровне LLM (то есть окончательной генерации текста с включением результатов top-k поиска), вы можете использовать:

bash scripts/run_target_attack.sh

Результаты и журналирование

• Журналы сохраняются в каталоге logs/.
• Результаты (например, отравленные запросы и показатели успешности атак) сохраняются в каталоге results/.

Благодарности

• Open-Prompt-Injection: Базовый код модели основан на этом репозитории.
• PoisonedRAG: Некоторые реализации были адаптированы из этого проекта.
• Beir Benchmark: Используется в качестве базовой платформы для оценки.
• Contriever: Используется для возможностей Retrieval-Augmented Generation (RAG).
• llms-attacks: Некоторые реализации для состязательного обучения были взяты из этого репозитория.