PoisonCraft

Este repositório fornece a implementação oficial do POISONCRAFT: Envenenamento Prático da Geração Aumentada por Recuperação para Grandes Modelos de Linguagem.

Visão Geral

POISONCRAFT tem como objetivo demonstrar como um agente malicioso pode plantar conteúdo "envenenado" no corpus usado por pipelines de Geração Aumentada por Recuperação (RAG), enganando um Grande Modelo de Linguagem para que ele alucine ou faça referência a conteúdo malicioso. Este código fornece scripts para:

• Preparar e pré-processar conjuntos de dados padrão (por exemplo, Natural Questions, MS MARCO, HotpotQA).
• Injetar sufixos adversariais (ou seja, "venenos") em texto semelhante a consultas, a fim de degradar ou manipular os resultados da recuperação.
• Avaliar a eficácia do envenenamento sob diferentes recuperadores (por exemplo, Contriever, SimCSE e BGE) e medir a transferibilidade.

Uso Rápido

Configuração do Ambiente

Abaixo está um guia de alto nível para configurar o ambiente. Os requisitos exatos estão listados em requirements.txt. Recomendamos criar um novo ambiente virtual para evitar conflitos:

# Exemplo usando conda
conda create -n PoisonCraft python=3.10
conda activate PoisonCraft

# Instalar dependências
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

Configurar Chave de API e Configuração do Modelo

Certas partes deste projeto, como classificação de domínio e categorização de consultas, dependem de chamadas de API (por exemplo, modelos GPT da OpenAI). Você precisa configurar a chave de API e as configurações do modelo da seguinte forma:

1. Navegue até o diretório src/models_configs/.
2. Localize o arquivo de configuração apropriado (por exemplo, gpt3.5_config.py).
3. Atualize o campo api_keys com sua chave de API:

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

Nota: Este projeto suporta vários provedores de LLM (por exemplo, OpenAI, Llama e PaLM2). Certifique-se de selecionar o arquivo de configuração apropriado para o seu modelo.

Como Executar o Código

Fornecemos vários scripts bash na pasta scripts para diferentes estágios do pipeline. Abaixo está uma referência passo a passo:

Preparando os Conjuntos de Dados

1. Baixar e Preparar Conjuntos de Dados

Execute scripts/run_prepare_dataset.sh para baixar e descompactar os conjuntos de dados, e criar divisões de treino/teste:

bash scripts/run_prepare_dataset.sh

Este script chama internamente:

• experiments/prepare_datasets.py
• experiments/classify_queries_by_domain.py (para classificação de domínio, se necessário)

2. Pré-processar Conjuntos de Dados

Para calcular as pontuações de similaridade ground truth top-k para cada domínio, execute:

bash scripts/run_process_data.sh

3. [Opcional] Avaliando BEIR para Ground Truth

Execute esta etapa apenas se quiser avaliar um modelo de recuperador, pois pode levar um tempo considerável. Siga estes passos:

1. Execute o script de avaliação BEIR para calcular os resultados da recuperação:

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

2. Use os resultados recuperados para gerar as pontuações ground truth:

   bash scripts/run_process_data.sh
   

Treinamento do Ataque de Envenenamento

Para gerar sufixos adversariais (ou seja, "venenos"), você pode usar:

bash scripts/run_multi_poisoncraft.sh

• Este script lançará várias execuções de run_poisoncraft.sh em paralelo, cada uma explorando diferentes comprimentos de sufixo adversarial e índices de domínio.
• A lógica principal chama experiments/poisoncraft_exp.py

Você pode personalizar parâmetros como:

• DOMAIN_LIST (os índices de domínio das consultas que deseja envenenar)
• ADV_LENGTH_LIST (os comprimentos dos sufixos adversariais a serem testados)

Executando a Avaliação do Ataque de Recuperação

Após gerar os sufixos adversariais, você pode avaliar quantas consultas obtêm referências "envenenadas" executando:

bash scripts/run_retrieval_attack.sh

Este script:

• Carrega os sufixos adversariais pré-computados.
• Executa um teste de ataque usando experiments/retrieval_attack.py.
• Registra os resultados de quantas consultas acabaram recuperando o texto envenenado.

Executando a Avaliação do Ataque de Geração Direcionada

Alternativamente, para avaliação no nível do LLM (ou seja, geração final de texto com resultados de recuperação top-k incluídos), você pode usar:

bash scripts/run_target_attack.sh

Resultados e Logging

• Os logs são armazenados no diretório logs/.
• Os resultados (por exemplo, consultas envenenadas e taxas de sucesso do ataque) são salvos no diretório results/.

Agradecimentos

• Open-Prompt-Injection: O código base do modelo é baseado neste repositório.
• PoisonedRAG: Algumas implementações foram adaptadas deste projeto.
• Beir Benchmark: Usado como estrutura de referência para avaliação.
• Contriever: Utilizado para capacidades de Geração Aumentada por Recuperação (RAG).
• llms-attacks: Certas implementações para treinamento adversarial foram derivadas deste repositório.