PoisonCraft

Este repositorio proporciona la implementación oficial de POISONCRAFT: Envenenamiento Práctico de la Generación Aumentada por Recuperación para Modelos de Lenguaje de Gran Tamaño.

Información General

POISONCRAFT tiene como objetivo demostrar cómo un actor malicioso puede plantar contenido “envenenado” en el corpus utilizado por los pipelines de Generación Aumentada por Recuperación (RAG), engañando así a un Modelo de Lenguaje de Gran Tamaño para que alucine o haga referencia a contenido malicioso. Este código base proporciona scripts para:

• Preparar y preprocesar conjuntos de datos estándar (p. ej., Natural Questions, MS MARCO, HotpotQA).

• Inyectar sufijos adversariales (es decir, “venenos”) en texto similar a consultas para degradar o manipular los resultados de recuperación.

• Evaluar la efectividad del envenenamiento con diferentes recuperadores (p. ej., Contriever, SimCSE y BGE) y medir la transferibilidad.

Uso Rápido

Configuración del Entorno

A continuación se presenta una guía de alto nivel para configurar el entorno. Los requisitos exactos se enumeran en requirements.txt. Recomendamos crear un nuevo entorno virtual para evitar conflictos:

# Example using conda
conda create -n PoisonCraft python=3.10
conda activate PoisonCraft

# Install dependencies
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 la Clave de API y la Configuración del Modelo

Ciertas partes de este proyecto, como la clasificación de dominios y la categorización de consultas, dependen de llamadas a la API (p. ej., modelos OpenAI GPT). Debe configurar la clave de API y los ajustes del modelo de la siguiente manera:

1. Navegue al directorio src/models_configs/.
2. Localice el archivo de configuración apropiado (p. ej., gpt3.5_config.py).
3. Actualice el campo api_keys con su clave de API:
   root@kitploit:~

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

Nota: Este proyecto admite múltiples proveedores de LLM (p. ej., OpenAI, Llama y PaLM2). Asegúrese de seleccionar el archivo de configuración apropiado para su modelo.

Cómo Ejecutar el Código

Proporcionamos varios scripts bash en la carpeta scripts para diferentes etapas del pipeline. A continuación se muestra una referencia paso a paso:

Preparación de los Conjuntos de Datos

1. Descargar y Preparar los Conjuntos de Datos

Ejecute scripts/run_prepare_dataset.sh para descargar y descomprimir los conjuntos de datos, y crear divisiones de entrenamiento/prueba:

bash scripts/run_prepare_dataset.sh

Este script internamente llama a:

• experiments/prepare_datasets.py
• experiments/classify_queries_by_domain.py (para la clasificación de dominios, si es necesario)

2. Preprocesar los Conjuntos de Datos

Para calcular las puntuaciones de similitud ground truth top-k para cada dominio, ejecute:

bash scripts/run_process_data.sh

3. [Opcional]Evaluando BEIR para Ground Truth

Ejecute este paso solo si desea evaluar un modelo recuperador, ya que puede llevar bastante tiempo. Siga estos pasos:

1. Ejecute el script de evaluación de BEIR para calcular los resultados de recuperación:
   root@kitploit:~

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

2. Use los resultados recuperados para generar las puntuaciones ground truth:
   root@kitploit:~

   bash scripts/run_process_data.sh
   

Entrenamiento del Ataque de Envenenamiento

Para generar sufijos adversariales (es decir, “venenos”), puede usar:

bash scripts/run_multi_poisoncraft.sh

• Este script lanzará múltiples ejecuciones de run_poisoncraft.sh en paralelo, cada una explorando diferentes longitudes de sufijos adversariales e índices de dominio.
• La lógica central llama a experiments/poisoncraft_exp.py

Puede personalizar parámetros como:

• DOMAIN_LIST (los índices de dominio de las consultas que desea envenenar)
• ADV_LENGTH_LIST (las longitudes de sufijo adversarial a probar)

Ejecución de la Evaluación del Ataque de Recuperación

Después de generar los sufijos adversariales, puede evaluar cuántas consultas obtienen referencias “envenenadas” ejecutando:

bash scripts/run_retrieval_attack.sh

Este script:

• Carga los sufijos adversariales precomputados.
• Ejecuta una prueba de ataque usando experiments/retrieval_attack.py.
• Registra los resultados de cuántas consultas terminaron recuperando el texto envenenado.

Ejecución de la Evaluación del Ataque de Generación Dirigida

Alternativamente, para la evaluación a nivel de LLM (es decir, generación de texto final con los resultados de recuperación top-k incluidos), puede usar:

bash scripts/run_target_attack.sh

Resultados y Registro

• Los registros se almacenan en el directorio logs/.
• Los resultados (p. ej., consultas envenenadas y tasas de éxito del ataque) se guardan en el directorio results/.

Agradecimientos

• Open-Prompt-Injection: El código del modelo fundacional se basa en este repositorio.
• PoisonedRAG: Algunas implementaciones fueron adaptadas de este proyecto.
• Beir Benchmark: Utilizado como marco de evaluación comparativa para la evaluación.
• Contriever: Utilizado para capacidades de Generación Aumentada por Recuperación (RAG).
• llms-attacks: Ciertas implementaciones para el entrenamiento adversarial se derivaron de este repositorio.