Generating arbitrary text: a Markov chain algorithm - The Go Programming Language

A chain consists of a prefix and a suffix. Each prefix is a set number of words, while a suffix is a single word. A prefix can have an arbitrary number of suffixes. To model this data, we use a map[string][]string. Each map key is a prefix (a string) and its values are lists of suffixes (a slice of strings, []string).

Here is the example table from the package comment as modeled by this data structure:

map[string][]string{
 " ": {"I"},
 " I": {"am"},
 "I am": {"a", "not"},
 "a free": {"man!"},
 "am a": {"free"},
 "am not": {"a"},
 "a number!": {"I"},
 "number! I": {"am"},
 "not a": {"number!"},
}

While each prefix consists of multiple words, we store prefixes in the map as a single string. It would seem more natural to store the prefix as a[]string, but we can't do this with a map because the key type of a map must implement equality (and slices do not).

Therefore, in most of our code we will model prefixes as a[]string and join the strings together with a space to generate the map key:

Prefix Map key
[]string{"", ""} " "
[]string{"", "I"} " I"
[]string{"I", "am"} "I am"

The Chain struct has two unexported fields (those that do not begin with an upper case character), and so we write aNewChain constructor function that initializes thechain map with make and sets theprefixLen field.

This is constructor function is not strictly necessary as this entire program is within a single package (main) and therefore there is little practical difference between exported and unexported fields. We could just as easily write out the contents of this function when we want to construct a new Chain. But using these unexported fields is good practice; it clearly denotes that only methods of Chain and its constructor function should access those fields. Also, structuring Chain like this means we could easily move it into its own package at some later date.

The Build method reads text from an io.Reader and parses it into prefixes and suffixes that are stored in theChain.

The io.Reader is an interface type that is widely used by the standard library and other Go code. Our code uses thefmt.Fscan function, which reads space-separated values from an io.Reader.

The Build method returns once the Reader'sRead method returns io.EOF (end of file) or some other read error occurs.

This function does many small reads, which can be inefficient for someReaders. For efficiency we wrap the providedio.Reader withbufio.NewReader to create a new io.Reader that provides buffering.

In our loop we read words from the Reader into astring variable s usingfmt.Fscan. Since Fscan uses space to separate each input value, each call will yield just one word (including punctuation), which is exactly what we need.

Fscan returns an error if it encounters a read error (io.EOF, for example) or if it can't scan the requested value (in our case, a single string). In either case we just want to stop scanning, so we break out of the loop.

The word stored in s is a new suffix. We add the new prefix/suffix combination to the chain map by computing the map key with p.String and appending the suffix to the slice stored under that key.

The built-in append function appends elements to a slice and allocates new storage when necessary. When the provided slice isnil, append allocates a new slice. This behavior conveniently ties in with the semantics of our map: retrieving an unset key returns the zero value of the value type and the zero value of []string is nil. When our program encounters a new prefix (yielding a nil value in the map) append will allocate a new slice.

For more information about the append function and slices in general see theSlices: usage and internals article.

Before reading the next word our algorithm requires us to drop the first word from the prefix and push the current suffix onto the prefix.

When in this state

p == Prefix{"I", "am"}
s == "not" 

the new value for p would be

p == Prefix{"am", "not"}

This operation is also required during text generation so we put the code to perform this mutation of the slice inside a method onPrefix named Shift.

The Shift method uses the built-in copy function to copy the last len(p)-1 elements of p to the start of the slice, effectively moving the elements one index to the left (if you consider zero as the leftmost index).

p := Prefix{"I", "am"}
copy(p, p[:1])
// p == Prefix{"am", "am"}

We then assign the provided word to the last index of the slice:

// suffix == "not"
p[len(p)-1] = suffix
// p == Prefix{"am", "not"}

The Generate method is similar to Build except that instead of reading words from a Reader and storing them in a map, it reads words from the map and appends them to a slice (words).

Generate uses a conditional for loop to generate up to n words.

At each iteration of the loop we retrieve a list of potential suffixes for the current prefix. We access the chain map at keyp.String() and assign its contents to choices.

If len(choices) is zero we break out of the loop as there are no potential suffixes for that prefix. This test also works if the key isn't present in the map at all: in that case, choices will be nil and the length of a nil slice is zero.

To choose a suffix we use therand.Intn function. It returns a random integer up to (but not including) the provided value. Passing in len(choices) gives us a random index into the full length of the list.

We use that index to pick our new suffix, assign it tonext and append it to the words slice.

Next, we Shift the new suffix onto the prefix just as we did in the Build method.

To make it easy to tweak the prefix and generated text lengths we use the flag package to parse command-line flags.

These calls to flag.Int register new flags with theflag package. The arguments to Int are the flag name, its default value, and a description. The Int function returns a pointer to an integer that will contain the user-supplied value (or the default value if the flag was omitted on the command-line).

The main function begins by parsing the command-line flags with flag.Parse and seeding the rand package's random number generator with the current time.

If the command-line flags provided by the user are invalid theflag.Parse function will print an informative usage message and terminate the program.

To create the new Chain we call NewChain with the value of the prefix flag.

To build the chain we call Build withos.Stdin (which implements io.Reader) so that it will read its input from standard input.

Finally, to generate text we call Generate with the value of the words flag and assigning the result to the variable text.

Then we call fmt.Println to write the text to standard output, followed by a carriage return.

To use this program, first build it with thego command:

$ go build markov.go

And then execute it while piping in some input text:

$ echo "a man a plan a canal panama" \
 | ./markov -prefix=1
a plan a man a plan a canal panama

Here's a transcript of generating some text using the Go distribution's README file as source material:

$ ./markov -words=10 < $GOROOT/README
This is the source code repository for the Go source
$ ./markov -prefix=1 -words=10 < $GOROOT/README
This is the go directory (the one containing this README).
$ ./markov -prefix=1 -words=10 < $GOROOT/README
This is the variable if you have just untarred a