vec2text

This library contains code for doing text embedding inversion. We can train various architectures that reconstruct text sequences from embeddings as well as run pre-trained models. This repository contains code for the papers "Text Embeddings Reveal (Almost) As Much As Text" and "Language Model Inversion".

To get started, install this on PyPI:

pip install vec2text

Link to Colab Demo

Development

If you're training a model you'll need to set up nltk:

import nltk
nltk.download('punkt')

Before pushing any code, please run precommit:

pre-commit run --all

Usage

The library can be used to embed text and then invert it, or invert directly from embeddings. First you'll need to construct a Corrector object which wraps the necessary models, embedders, and tokenizers:

Load a model via load_pretrained_corrector

corrector = vec2text.load_pretrained_corrector("text-embedding-ada-002")

Load a model via load_corrector

If you have trained you own custom models using vec2text, you can load them in using the load_corrector function.

inversion_model = vec2text.models.InversionModel.from_pretrained("jxm/gtr__nq__32")
corrector_model = vec2text.models.CorrectorEncoderModel.from_pretrained("jxm/gtr__nq__32__correct")

corrector = vec2text.load_corrector(inversion_model, corrector_model)

Both vec2text.models.InversionModel and vec2text.models.CorrectorEncoderModel classes inherit transformers.PreTrainedModel therefore you can pass in a Hugging Face model name or path to a local directory.

Invert text with invert_strings

vec2text.invert_strings(
    [
        "Jack Morris is a PhD student at Cornell Tech in New York City",
        "It was the best of times, it was the worst of times, it was the age of wisdom, it was the age of foolishness, it was the epoch of belief, it was the epoch of incredulity"
    ],
    corrector=corrector,
)
['Morris is a PhD student at Cornell University in New York City',
 'It was the age of incredulity, the age of wisdom, the age of apocalypse, the age of apocalypse, it was the age of faith, the age of best faith, it was the age of foolishness']

By default, this will make a single guess (using the hypothesizer). For better results, you can make multiple steps:

vec2text.invert_strings(
    [
        "Jack Morris is a PhD student at Cornell Tech in New York City",
        "It was the best of times, it was the worst of times, it was the age of wisdom, it was the age of foolishness, it was the epoch of belief, it was the epoch of incredulity"
    ],
    corrector=corrector,
    num_steps=20,
)
['Jack Morris is a PhD student in tech at Cornell University in New York City',
 'It was the best time of the epoch, it was the worst time of the epoch, it was the best time of the age of wisdom, it was the age of incredulity, it was the age of betrayal']

And for even better results, you can increase the size of the search space by setting sequence_beam_width to a positive integer:

vec2text.invert_strings(
    [
        "Jack Morris is a PhD student at Cornell Tech in New York City",
        "It was the best of times, it was the worst of times, it was the age of wisdom, it was the age of foolishness, it was the epoch of belief, it was the epoch of incredulity"
    ],
    corrector=corrector,
    num_steps=20,
    sequence_beam_width=4,
)
['Jack Morris is a PhD student at Cornell Tech in New York City',
 'It was the best of times, it was the worst of times, it was the age of wisdom, it was the age of foolishness, it was the epoch of belief, it was the epoch of incredulity']

Note that this technique has to store sequence_beam_width * sequence_beam_width hypotheses at each step, so