python experiments/preprocess_data.py
--model_name_or_path lmsys/vicuna-7b-v1.3
--raw_dataset_path /path/to/alpaca_data_cleaned.json
--output_dir <output_dir>
--num_samples 500 # number of prompts to include in the dataset
--max_len 32 # number of tokens in each generated document
--num_docs_per_sample 100 # total number of docs for each prompt
--batch_size 100 # make this as high as possible without causing CUDA OOM

python experiments/run_experiments_soft.py
--model_name_or_path lmsys/vicuna-7b-v1.3
--dataset_path /path/to/generated_datasets.pkl # path to the dataset you generated (NOT the raw data)
--output_dir <output_dir>
--batch_size 10
--learning_rate 3e-4
--num_epochs 50
--early_stopping # whether to stop once KL starts increasing again
--kl_every 5 # how often to compute KL (it slows it down slightly if its small)
--optim_suffix_len 10
--n_trials 10

python experiments/run_experiments_hard.py
--model_name_or_path lmsys/vicuna-7b-v1.3
--dataset_path /path/to/generated_datasets.pkl # path to the dataset you generated (NOT the raw data)
--output_dir <output_dir>
--batch_size 10
--top_k 250 # number of tokens to consider for top gradients
--num_epochs 50
--n_proposals # number of proposals to generate per iter. If its higher than the total number of tokens, it'll default to the total # of tokens (ofc for best convergence, use all the tokens)
--kl_every 5
--optim_suffix_len 10 # length of prompt to optimize... anywhere from 10-20 seems fine without taking too much GPU mem
--init_prompt_char # character to initialize prompt with, we just use the default ! that llm-attacks used
--clip_vocab # prune the tokenizer's vocabulary to only consider English tokens
--warm_start_file # path to file with the chatgpt suggested prompts to uniformly sample for warm starts
--n_trials 10
--sharded # use model parallel (tested on 4x 4090, we can't fit the optimization in 24GB so we need 2 GPUs per training)

@article{melamed2023propane,
  title={PROPANE: Prompt design as an inverse problem},
  author={Melamed, Rimon and McCabe, Lucas H and Wakhare, Tanay and Kim, Yejin and Huang, H Howie and Boix-Adsera, Enric},
  journal={arXiv preprint arXiv:2311.07064},
  year={2023}
}