This repository contains code necessary to replicate the training and evaluation for our ICLR 2023 paper "Systematic Rectification of Language Models via Dead-end Analysis" by Meng Cao, Mehdi Fatemi, Jackie CK Cheung and Samira Shabanian.
- Python >= 3.8
- PyTorch >= 1.7.1
- transformers >= 4.22.0
- accelerate >= 0.12.0
- scikit-learn >= 1.0.2
Begin by creating a training dataset according to the guidelines provided in the dataset folder. Alternatively, you can download our pre-processed training data from Google drive. After obtaining the train.json
and val.json
files, execute the command below to start training:
OUTPUT_DIR=checkpoint/
TRAIN_FILE=dataset/qd_dataset/processed/train.json
VALID_FILE=dataset/qd_dataset/processed/val.json
accelerate launch --config_file accelerator_config.yaml src/train.py \
--overwrite_cache true \
--gamma 1.0 \
--num_train_epochs 3 \
--q_max 0.0 --q_min -1.0 \
--per_device_train_batch_size 64 \
--per_device_eval_batch_size 64 \
--preprocessing_num_workers 16 \
--num_warmup_steps 500 \
--polyak_update_lr 0.1 \
--gradient_accumulation_steps 1 \
--learning_rate 3e-4 \
--train_file $TRAIN_FILE \
--validation_file $VALID_FILE \
--model_name_or_path gpt2 \
--output_dir $OUTPUT_DIR \
--checkpointing_steps "epoch"
Our checkpoint can be found here.
We use Real Toxicity Prompts and Perspective API for evaluating language model toxicity. To access the API, follow the instructions here. You can download prompts-nontoxic-10k.jsonl
from here.
Q_PATH=checkpoint/
SAVE_PATH=generations.jsonl
python src/inference.py \
--lm_name_or_path "gpt2-large" \
--q_model_path $Q_PATH \
--prompts_path dataset/prompts-nontoxic-10k.jsonl \
--save_path $SAVE_PATH \
--do_sample \
--batch_size 3 \
--num_samples_for_eval 5000 \
--epsilon 0.1;
echo "Run Perspective API for toxicity evaluation..."
python evaluation/run_perspective_api.py $SAVE_PATH \
--num_return 25 \
--num_thread 8 \
--perspective_api_key $PERSPECTIVE_API_KEY \
--save_scores --score_saving_path "eval_scores.jsonl";
echo "Diversity evaluation..."
python evaluation/diversity_eval.py -f $SAVE_PATH --num_return 25
Please cite as:
@inproceedings{
cao2023systematic,
title={Systematic Rectification of Language Models via Dead-end Analysis},
author={Meng Cao and Mehdi Fatemi and Jackie CK Cheung and Samira Shabanian},
booktitle={The Eleventh International Conference on Learning Representations },
year={2023},
url={https://openreview.net/forum?id=k8_yVW3Wqln}
}