This is the official repo for the paper "Quark: Controllable Text Generation with Reinforced Unlearning" (NeurIPS 2022)

We suggest using conda to setup environment. You need to first replace prefix in environment.yml with your home path. With conda installed, create an environment called quark with:

conda env create -f environment.yml

The main branch contains toxicity unlearning task. We put the other two tasks, sentiment steering and repetition reduction in sentiment branch and repetition branch separately.

We use the PerspectiveAPI to score toxicity in reward computing, which requires API key for access. Please refer to their website for API key application.

Please first replace PERSPECTIVE_API_KEY in constants.py with your own API key. For training quark for toxicity reduction with default hyperparameters,

python main.py

You can change hyperparameters in arguments.py via argparse.

To evaluate the toxicity of unlearned model, please use sample.py. You need to first replace save_path and checkpoint_path with your output directory and model checkpoint path, then

python sample.py

It will save the evaluation result to your output directory.

To evaluate perplexity of the generations, please use perplexity.py. You need to first replace save_path with the same output directory specified above, then

python perplexity.py

It will save the perplexity result to the same output directory.

We release our model checkpoints for all three tasks: toxicity unlearn, sentiment steering (positive, negative) and repetition reduction.

If you use this codebase in your work, please consider citing our paper:

@article{Lu2022QuarkCT,
  title={Quark: Controllable Text Generation with Reinforced Unlearning},
  author={Ximing Lu and Sean Welleck and Liwei Jiang and Jack Hessel and Lianhui Qin and Peter West and Prithviraj Ammanabrolu and Yejin Choi},
  journal={ArXiv},
  year={2022},
  volume={abs/2205.13636}
}