DIStributional COntrol

🕺🏽 disco allows you to flexibly control the generation from language models and other generative models.

Get it with pip install disco-generation.

Try it now!

example.py | constrain GPT-2 to always include "amazing"

# load a language model
distribution = LMDistribution("gpt2")

# define the feature we want to control for:
# "amazing" appears in the generated text
amazing = BooleanScorer(lambda s, c: "amazing" in s.text)

# define our target distribution:
# the feature amazing should appear 100% of the time
target = base.constrain([amazing], [1.0])

# create a model to fine-tune
model = LMDistribution("gpt2", freeze=False)

# fine-tune the model to approximate the target distribution
tuner = DPGTuner(model, target)
tuner.tune()

# generate amazing texts
samples, log_scores = model.sample()
for s in samples:
  print(s.text)