рекламный суфлер

AdvPrompter: быстрые адаптивные состязательные подсказки для LLM

Этот репозиторий является официальной реализацией AdvPrompter (arxiv:2404.16873).

Пожалуйста? звезда? это репозиторий и цитируйте нашу статью, если вам нравится (и/или вы используете) нашу работу, спасибо!

0. Установка

• Вариант 1. Используйте Singularity для установки контейнера advprompter.def.
• Вариант 2. Установите python3.11 и require.txt вручную:

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

1. Запуск AdvPrompter

Мы используем гидру в качестве инструмента управления конфигурацией. Основные файлы конфигурации: ./conf/{train,eval,eval_suffix_dataset,base}.yaml . AdvPrompter и TargetLLM указаны в conf/base.yaml, различные параметры уже реализованы.

База кода дополнительно поддерживает wandb, установив соответствующие параметры в conf/base.yaml.

Обратите внимание на характеристики оборудования: во всех проведенных нами экспериментах использовались два графических процессора NVIDIA A100: один для AdvPrompter и один для TargetLLM. Вы можете управлять устройствами в conf/target_llm/base_target_llm.yaml и conf/prompter/base_prompter.yaml.

1.1 Оценка

Бегать

python3 main.py --config-name=eval

для проверки производительности указанного AdvPrompter по сравнению с TargetLLM в заданном наборе данных. Вам нужно будет указать TargetLLM и AdvPrompter в conf/base.yaml. Кроме того, вы можете указать путь к peft_checkpoint, если AdvPrompter был настроен ранее:

// см. conf/prompter/llama2.yaml
лора_парамс:
теплый старт: правда
lora_checkpoint: "path_to_peft_checkpoint"

Суффиксы, созданные во время оценки, сохраняются в новом наборе данных в каталоге запуска в ./exp/.../suffix_dataset для дальнейшего использования. Такой набор данных также может быть полезен для оценки базовых показателей или созданных вручную суффиксов по сравнению с TargetLLM, и его можно оценить, запустив

python3 main.py --config-name=eval_suffix_dataset

после заполнения suffix_dataset_pth_dct в eval_suffix_dataset.yaml

1.2. Обучение

Бегать

python3 main.py --config-name=train

для обучения указанного AdvPrompter работе с TargetLLM. Он автоматически выполняет указанную выше оценку через регулярные промежутки времени, а также сохраняет промежуточные версии AdvPrompter в каталоге запуска под ./exp/.../checkpoints для последующего теплого запуска. Контрольную точку можно указать с помощью параметра lora_checkpoint в конфигурациях модели (как показано в разделе 1.1 «Оценка»). Обучение также сохраняет для каждой эпохи целевые суффиксы, сгенерированные с помощью AdvPrompterOpt, в ./exp/.../suffix_opt_dataset . Это позволяет предварительно обучать такой набор данных суффиксов, указав соответствующий путь в разделе pretrain в train.yaml

Некоторые важные гиперпараметры, которые следует учитывать в conf/train.yaml: [epochs, lr, top_k, num_chunks, lambda_val]

Примеры

Примечание. Возможно, вы захотите заменить target_llm.llm_params.checkpoint локальным путем.

• Пример 1: AdvPrompter на Vicuna-7B:

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

• Пример 2: AdvPrompter на 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

• Пример 3: AdvPrompter в чате Мистраль-7Б:

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

• Пример 4: AdvPrompter в чате Llama2-7B:

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

2. Участники

Ансельм Паулюс*, Арман Жармагамбетов*, Чуан Го, Брэндон Амос**, Юаньдун Тянь**

(* = равные первые авторы, ** = равные консультации)

3. Лицензия

Наш исходный код находится под лицензией CC-BY-NC 4.0.