Getting Started

How about we start with a basic algorithm first for watermarking? This follows the concept of green list and red list tokens.

Let me explain a bit more, when the model is predicting the next token, it is not allowed to choose any of the tokens from the red list as the next token.

How do we divide the tokens into red list and green list tokens you might ask?

Let's make and run our own very basic simulation for this, First up, let's define some variables that we really need which is our model's vocabulary and the prompt that we will send!

vocabulary = ['the', 'quick', 'brown', 'fox', 'jumps', 'over', 'lazy', 'dog', 'runs', 'and']
prompt = "The quick brown"

Now, let's define our language_model function which will just assign random logits to each word for now and return the softmax probability vector for the purpose of the simulation.

def language_model(vocabulary):
    # Assign random logits to each word in the vocabulary
    logits = np.random.randn(len(vocabulary))
 
    # Convert logits to probabilities using softmax
    probabilities = np.exp(logits) / np.sum(np.exp(logits))
 
    return dict(zip(vocabulary, probabilities))

Try running this cell below and see what it prints, you can even tinker with this code cell if you want

# Apply the language model to get the probability vector
prob_vector = language_model(vocabulary)
print(prob_vector)

Now, we need to compute a hash. In simple terms, a hash is like a fingerprint for data. It is a unique string of characters generated by a mathematical algorithm from any input data. This string is typically of a fixed length, regardless of the size of the input. If the input to the hash function changes, the hash output also changes, but the same input will always give the same output.
Try it out for yourself 👇

import hashlib
token='brown'
print(hashlib.sha256(token.encode()).hexdigest())

Phew, a lot of work done already Now let's use the value generated by the hash above to seed a random number generator

random.seed(hash_value)

The generator is seeded with the last token, which enables the red list to be reproduced later without needing access to the entire generated text. Now we will use the random number generator to partition the vocabulary into red list and green list tokens!

random.shuffle(vocabulary)
half_size = len(vocabulary) // 2
green_list = vocabulary[:half_size]
red_list = vocabulary[half_size:]
print("red list: ", red_list)
print("green list: ", green_list)

We will sample the next token from the green list, never generating any token in the red list. For now, let's pick the token from the green list with the highest probability. Run this and see what you get!

next_token = max(green_list, key=lambda word: prob_vector[word])
print(next_token)

Cool, I think the simulation is pretty much complete. In the next section I have the full algorithm coded out.