PoisonCraft

Dieses Repository enthält die offizielle Implementierung von POISONCRAFT: Practical Poisoning of Retrieval-Augmented Generation for Large Language Models.

Überblick

POISONCRAFT zielt darauf ab zu demonstrieren, wie ein böswilliger Akteur „vergiftete" Inhalte in den Korpus einschleusen kann, der von Retrieval-Augmented-Generation-Pipelines (RAG) verwendet wird, und so ein Large Language Model dazu bringt, zu halluzinieren oder auf schädliche Inhalte zu verweisen. Diese Codebasis stellt Skripte bereit, um:

• Standarddatensätze (z. B. Natural Questions, MS MARCO, HotpotQA) vorzubereiten und zu verarbeiten.

• Adversariale Suffixe (d. h. „Gifte") in query-ähnlichen Text einzufügen, um Retrieval-Ergebnisse zu verschlechtern oder zu manipulieren.

• Die Wirksamkeit der Vergiftung unter verschiedenen Retrievern (z. B. Contriever, SimCSE und BGE) zu bewerten und die Übertragbarkeit zu messen.

Schnelle Nutzung

Umgebungseinrichtung

Im Folgenden finden Sie eine allgemeine Anleitung zur Konfiguration der Umgebung. Die genauen Anforderungen sind in requirements.txt aufgeführt. Wir empfehlen, eine neue virtuelle Umgebung zu erstellen, um Konflikte zu vermeiden:

# Beispiel mit conda
conda create -n PoisonCraft python=3.10
conda activate PoisonCraft

# Abhängigkeiten installieren
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-Schlüssel und Modellkonfiguration festlegen

Bestimmte Teile dieses Projekts, wie z. B. die Domänenklassifizierung und die Query-Kategorisierung, sind auf API-Aufrufe angewiesen (z. B. OpenAI-GPT-Modelle). Sie müssen den API-Schlüssel und die Modelleinstellungen wie folgt konfigurieren:

1. Navigieren Sie zum Verzeichnis src/models_configs/.
2. Suchen Sie die passende Konfigurationsdatei (z. B. gpt3.5_config.py).
3. Aktualisieren Sie das Feld api_keys mit Ihrem API-Schlüssel:
   root@kitploit:~

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

Hinweis: Dieses Projekt unterstützt mehrere LLM-Anbieter (z. B. OpenAI, Llama und PaLM2). Stellen Sie sicher, dass Sie die passende Konfigurationsdatei für Ihr Modell auswählen.

So führen Sie den Code aus

Wir stellen verschiedene Bash-Skripte im Ordner scripts für die verschiedenen Phasen der Pipeline bereit. Im Folgenden finden Sie eine Schritt-für-Schritt-Referenz:

Vorbereiten der Datensätze

1. Datensätze herunterladen und vorbereiten

Führen Sie scripts/run_prepare_dataset.sh aus, um die Datensätze herunterzuladen und zu entpacken sowie Trainings-/Test-Splits zu erstellen:

bash scripts/run_prepare_dataset.sh

Dieses Skript ruft intern Folgendes auf:

• experiments/prepare_datasets.py
• experiments/classify_queries_by_domain.py (für die Domänenklassifizierung, falls erforderlich)

2. Datensätze vorverarbeiten

Um die Top-k-Ground-Truth-Ähnlichkeitswerte für jede Domäne zu berechnen, führen Sie Folgendes aus:

bash scripts/run_process_data.sh

3. [Optional] BEIR für Ground Truth auswerten

Führen Sie diesen Schritt nur aus, wenn Sie ein Retriever-Modell bewerten möchten, da dies erheblich Zeit in Anspruch nehmen kann. Gehen Sie wie folgt vor:

1. Führen Sie das BEIR-Bewertungsskript aus, um die Retrieval-Ergebnisse zu berechnen:
   root@kitploit:~

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

2. Verwenden Sie die abgerufenen Ergebnisse, um Ground-Truth-Scores zu erzeugen:
   root@kitploit:~

   bash scripts/run_process_data.sh
   

Training des Vergiftungsangriffs

Um adversariale Suffixe (d. h. „Gifte") zu erzeugen, können Sie Folgendes verwenden:

bash scripts/run_multi_poisoncraft.sh

• Dieses Skript startet mehrere Läufe von run_poisoncraft.sh parallel, wobei jeweils unterschiedliche Längen der adversarialen Suffixe und Domänenindizes untersucht werden.
• Die Kernlogik ruft experiments/poisoncraft_exp.py auf.

Sie können Parameter wie folgt anpassen:

• DOMAIN_LIST (die Domänenindizes der Queries, die Sie vergiften möchten)
• ADV_LENGTH_LIST (die zu testenden Längen der adversarialen Suffixe)

Auswertung des Retrieval-Angriffs ausführen

Nach dem Erzeugen der adversarialen Suffixe können Sie auswerten, wie viele Queries „vergiftete" Referenzen erhalten, indem Sie Folgendes ausführen:

bash scripts/run_retrieval_attack.sh

Dieses Skript:

• Lädt die vorab berechneten adversarialen Suffixe.
• Führt einen Angriffstest mit experiments/retrieval_attack.py durch.
• Protokolliert die Ergebnisse, wie viele Queries am Ende den vergifteten Text abgerufen haben.

Auswertung des gezielten Generierungsangriffs ausführen

Alternativ können Sie für die Bewertung auf LLM-Ebene (d. h. die finale Textgenerierung einschließlich der Top-k-Retrieval-Ergebnisse) Folgendes verwenden:

bash scripts/run_target_attack.sh

Ergebnisse und Protokollierung

• Protokolle werden im Verzeichnis logs/ gespeichert.
• Ergebnisse (z. B. vergiftete Queries und Angriffserfolgsraten) werden im Verzeichnis results/ gespeichert.

Danksagung

• Open-Prompt-Injection: Der grundlegende Modellcode basiert auf diesem Repository.
• PoisonedRAG: Einige Implementierungen wurden aus diesem Projekt übernommen.
• Beir Benchmark: Wird als Benchmark-Framework für die Bewertung verwendet.
• Contriever: Wird für Retrieval-Augmented-Generation-Fähigkeiten (RAG) genutzt.
• llms-attacks: Bestimmte Implementierungen für adversariales Training wurden aus diesem Repository abgeleitet.