advprompter

AdvPrompter: Solicitação Adversarial Adaptativa Rápida para LLMs

Este repositório é uma implementação oficial do AdvPrompter (arxiv:2404.16873).

Por favor? Estrela? este repositório e cite nosso artigo se você gosta (e/ou usa) nosso trabalho, obrigado!

0. Instalação

• Opção 1. Use Singularity para instalar o contêiner advprompter.def
• Opção 2. Instale python3.11 e requisitos.txt manualmente:

  conda create -n advprompter python=3.11.4
  conda activate advprompter
  pip install -r requirements.txt

1. Executando AdvPrompter

Usamos Hydra como ferramenta de gerenciamento de configuração. Arquivos de configuração principais: ./conf/{train,eval,eval_suffix_dataset,base}.yaml O AdvPrompter e o TargetLLM são especificados em conf/base.yaml, várias opções já estão implementadas.

A base de código suporta opcionalmente wandb definindo as opções correspondentes em conf/base.yaml.

Nota sobre as especificações de hardware: todos os experimentos que conduzimos utilizaram duas GPUs NVIDIA A100, uma para AdvPrompter e outra para TargetLLM. Você pode gerenciar dispositivos em conf/target_llm/base_target_llm.yaml e conf/prompter/base_prompter.yaml

1.1 Avaliação

Correr

python3 main.py --config-name=eval

para testar o desempenho do AdvPrompter especificado em relação ao TargetLLM em um determinado conjunto de dados. Você terá que especificar TargetLLM e AdvPrompter em conf/base.yaml. Além disso, você pode querer especificar um caminho para peft_checkpoint se o AdvPrompter tiver sido ajustado antes:

// veja conf/prompter/llama2.yaml
lora_params:
início quente: verdadeiro
lora_checkpoint: "caminho_para_peft_checkpoint"

Os sufixos gerados durante a avaliação são salvos em um novo conjunto de dados no diretório de execução em ./exp/.../suffix_dataset para uso posterior. Esse conjunto de dados também pode ser útil para avaliar linhas de base ou sufixos artesanais em relação a um TargetLLM e pode ser avaliado executando

python3 main.py --config-name=eval_suffix_dataset

depois de preencher o suffix_dataset_pth_dct em eval_suffix_dataset.yaml

1.2. Treinamento

Correr

python3 main.py --config-name=train

para treinar o AdvPrompter especificado em relação ao TargetLLM. Ele executa automaticamente a avaliação especificada acima em intervalos regulares e também salva versões intermediárias do AdvPrompter no diretório de execução em ./exp/.../checkpoints para início a quente posterior. O ponto de verificação pode ser especificado com o parâmetro lora_checkpoint nas configurações do modelo (conforme ilustrado na Avaliação 1.1). O treinamento também salva para cada época os sufixos de destino gerados com AdvPrompterOpt em ./exp/.../suffix_opt_dataset . Isso permite o pré-treinamento em tal conjunto de dados de sufixos, especificando o caminho correspondente em pré-treinamento em train.yaml

Alguns hiperparâmetros importantes a serem considerados em conf/train.yaml: [epochs, lr, top_k, num_chunks, lambda_val]

Exemplos

Nota: você pode querer substituir target_llm.llm_params.checkpoint por um caminho local.

• Exemplo 1: AdvPrompter em Vicuna-7B:

   python3 main.py --config-name=train target_llm=vicuna_chat target_llm.llm_params.model_name=vicuna-7b-v1.5

• Exemplo 2: AdvPrompter em Vicuna-13B:

   python3 main.py --config-name=train target_llm=vicuna_chat target_llm.llm_params.model_name=vicuna-13b-v1.5 target_llm.llm_params.checkpoint=lmsys/vicuna-13b-v1.5 train.q_params.num_chunks=2

• Exemplo 3: AdvPrompter no chat Mistral-7B:

   python3 main.py --config-name=train target_llm=mistral_chat

• Exemplo 4: AdvPrompter no chat Llama2-7B:

   python3 main.py --config-name=train target_llm=llama2_chat train.q_params.lambda_val=150

2. Contribuintes

Anselmo Paulus*, Arman Zharmagambetov*, Chuan Guo, Brandon Amos**, Yuandong Tian**

(* = Igualdade de 1ºs autores, ** = Igualdade de orientação)

3. Licença

Nosso código-fonte está sob licença CC-BY-NC 4.0.