liamjxu/mucoco

Official Code for the papers: "Controlled Text Generation as Continuous Optimization with Multiple Constraints" and "Gradient-based Constrained Sampling from LMs"

This repository contains the code for the NeurIPS 2021 paper: Controlled Text Generation as Continuous Optimization with Multiple Constraints

Dependencies

• pytorch >= 1.6
• transformers >= 4.5.1
• (optional for some constraints) sentence-transformers
• (optional for some constraints) POT

Quick Start

The main file to run this decoding algorithm is decode.py. All models used in this code are based on huggingface transformers.

Machine Translation experiments

see examples

Style Transfer experiments

see examples

Adding new constraints

This code currently supports the following losses:

• Sentence Classification (Cross Entropy)
• Semantic Similarity (Cosine Similarity, WMD between representations)
• Conditional generation losses (MarianMT, GPT2)

To add more losses/constraints, follow examples from 'mucoco/losses/'

License

The source code is licensed under the MIT license, which you can find in the LICENSE.md file