PoisonCraft

Questo repository fornisce l'implementazione ufficiale di POISONCRAFT: Practical Poisoning of Retrieval-Augmented Generation for Large Language Models.

Panoramica

POISONCRAFT mira a dimostrare come un attore malintenzionato possa inserire contenuti "avvelenati" nel corpus utilizzato dalle pipeline di Retrieval-Augmented Generation (RAG), inducendo così un Large Language Model ad allucinare o a fare riferimento a contenuti dannosi. Questo codebase fornisce script per:

• Preparare e preprocessare dataset standard (ad es., Natural Questions, MS MARCO, HotpotQA).

• Iniettare suffissi avversari (cioè "veleni") in testi simili a query al fine di degradare o manipolare i risultati del retrieval.

• Valutare l'efficacia dell'avvelenamento con diversi retriever (ad es., Contriever, SimCSE e BGE) e misurarne la trasferibilità.

Utilizzo Rapido

Configurazione dell'Ambiente

Di seguito una guida di alto livello per configurare l'ambiente. I requisiti esatti sono elencati in requirements.txt. Consigliamo di creare un nuovo ambiente virtuale per evitare conflitti:

# Esempio con conda
conda create -n PoisonCraft python=3.10
conda activate PoisonCraft

# Installa le dipendenze
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

Impostare la Chiave API e la Configurazione del Modello

Alcune parti di questo progetto, come la classificazione per dominio e la categorizzazione delle query, dipendono da chiamate API (ad es., modelli OpenAI GPT). È necessario configurare la chiave API e le impostazioni del modello come segue:

1. Navigare nella directory src/models_configs/.
2. Individuare il file di configurazione appropriato (ad es., gpt3.5_config.py).
3. Aggiornare il campo api_keys con la propria chiave API:
   root@kitploit:~

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

Nota: Questo progetto supporta più provider LLM (ad es., OpenAI, Llama e PaLM2). Assicurarsi di selezionare il file di configurazione appropriato per il proprio modello.

Come Eseguire il Codice

Forniamo vari script bash nella cartella scripts per le diverse fasi della pipeline. Di seguito una guida passo passo:

Preparazione dei Dataset

1. Scaricare e Preparare i Dataset

Eseguire scripts/run_prepare_dataset.sh per scaricare e decomprimere i dataset e creare le suddivisioni train/test:

bash scripts/run_prepare_dataset.sh

Questo script richiama internamente:

• experiments/prepare_datasets.py
• experiments/classify_queries_by_domain.py (per la classificazione per dominio, se necessario)

2. Preprocessare i Dataset

Per calcolare i punteggi di similarità top-k rispetto alla ground truth per ciascun dominio, eseguire:

bash scripts/run_process_data.sh

3. [Opzionale] Valutazione BEIR per la Ground Truth

Eseguire questo passaggio solo se si desidera valutare un modello retriever, poiché potrebbe richiedere molto tempo. Seguire questi passaggi:

1. Eseguire lo script di valutazione BEIR per calcolare i risultati del retrieval:
   root@kitploit:~

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

2. Utilizzare i risultati recuperati per generare i punteggi ground truth:
   root@kitploit:~

   bash scripts/run_process_data.sh
   

Addestramento dell'Attacco di Avvelenamento

Per generare i suffissi avversari (cioè i "veleni"), si può usare:

bash scripts/run_multi_poisoncraft.sh

• Questo script eseguirà più run di run_poisoncraft.sh in parallelo, ciascuna esplorando diverse lunghezze di suffissi avversari e indici di dominio.
• La logica principale richiama experiments/poisoncraft_exp.py

È possibile personalizzare parametri come:

• DOMAIN_LIST (gli indici di dominio delle query che si desidera avvelenare)
• ADV_LENGTH_LIST (le lunghezze dei suffissi avversari da testare)

Esecuzione della Valutazione dell'Attacco al Retrieval

Dopo aver generato i suffissi avversari, è possibile valutare quante query ricevono riferimenti "avvelenati" eseguendo:

bash scripts/run_retrieval_attack.sh

Questo script:

• Carica i suffissi avversari pre-calcolati.
• Esegue un test di attacco usando experiments/retrieval_attack.py.
• Registra i risultati relativi a quante query hanno finito per recuperare il testo avvelenato.

Esecuzione della Valutazione dell'Attacco di Generazione Mirata

In alternativa, per la valutazione a livello LLM (cioè la generazione finale del testo con l'inclusione dei risultati top-k del retrieval), si può usare:

bash scripts/run_target_attack.sh

Risultati e Registri

• I log sono salvati nella directory logs/.
• I risultati (ad es., query avvelenate e tassi di successo dell'attacco) sono salvati nella directory results/.

Ringraziamenti

• Open-Prompt-Injection: Il codice modello di base è basato su questo repository.
• PoisonedRAG: Alcune implementazioni sono state adattate da questo progetto.
• Beir Benchmark: Utilizzato come framework di benchmark per la valutazione.
• Contriever: Utilizzato per le capacità di Retrieval-Augmented Generation (RAG).
• llms-attacks: Alcune implementazioni per l'addestramento avversario derivano da questo repository.