PoisonCraft

このリポジトリは、POISONCRAFT: Practical Poisoning of Retrieval-Augmented Generation for Large Language Models の公式実装を提供します。

概要

POISONCRAFT は、悪意のある攻撃者が Retrieval-Augmented Generation（RAG）パイプラインで使用されるコーパスに「毒（poison）」を仕込んだコンテンツを埋め込み、大規模言語モデルを誤った情報の生成（ハルシネーション）や悪意のあるコンテンツの参照に誘導できることを実証することを目的としています。このコードベースは、以下のスクリプトを提供します。

• 標準データセット（例: Natural Questions、MS MARCO、HotpotQA）の準備と前処理。

• クエリ風テキストに敵対的サフィックス（すなわち「毒（poisons）」）を注入し、検索結果を劣化または操作する。

• 異なるレトリーバー（例: 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）をサポートしています。モデルに適した設定ファイルを選択してください。

コードの実行方法

パイプラインのさまざまな段階向けに、scripts フォルダの下にさまざまなbashスクリプトを提供しています。以下はステップバイステップのリファレンスです:

データセットの準備

1. データセットのダウンロードと準備

scripts/run_prepare_dataset.sh を実行して、データセットのダウンロードと解凍、およびトレーニング/テスト分割を作成します:

bash scripts/run_prepare_dataset.sh

このスクリプトは内部で以下を呼び出します:

• experiments/prepare_datasets.py
• experiments/classify_queries_by_domain.py（必要な場合のドメイン分類用）

2. データセットの前処理

各ドメインの上位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
   

ポイズニング攻撃トレーニング

敵対的サフィックス（すなわち「毒（poisons）」）を生成するには、次を使用できます:

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レベルの評価（すなわち、上位k件の検索結果を含む最終テキスト生成）には、次を使用できます:

bash scripts/run_target_attack.sh

結果とログ

• ログは logs/ ディレクトリに保存されます。
• 結果（例: 毒されたクエリと攻撃成功率）は results/ ディレクトリに保存されます。

謝辞

• Open-Prompt-Injection: 基礎となるモデルコードはこのリポジトリに基づいています。
• PoisonedRAG: 一部の実装はこのプロジェクトから適応されています。
• Beir Benchmark: 評価のベンチマークフレームワークとして使用。
• Contriever: Retrieval-Augmented Generation（RAG）機能のために利用。
• llms-attacks: 敵対的トレーニングの一部の実装はこのリポジトリから派生しています。