Vanilla Markov chain in 18 lines

Question

I am looking for advice on how I can improve this code. The objective is to make a vanilla Markov Chain text generator in the most concise way. I'm looking for advice on style, naming, 'elegance', and correctness.

import random, re

with open("corpus.txt") as file:
    text = file.read().replace('\n', ' newline ')

text = re.sub(r'([\,\.\!\?\;\:])',r' \1 ',text)
text = re.sub(r'[^a-zA-Z0-9|\,|\.|\!|\?|\;|\:]',r' ',text)

text = [word.lower().strip() for word in text.split(' ') if word.strip() != '']

markov_chain = {}
for current_word, next_word in zip(text, text[1:]):
    markov_chain.setdefault(current_word, []).append(next_word)

current_word = random.choice(list(markov_chain.keys()))
sentence = []

for i in range(0,15):
    sentence.append(current_word)
    if current_word in markov_chain and len(markov_chain[current_word]) > 0:
        current_word = random.choice(markov_chain[current_word])
    else:
        break

print(' '.join(sentence).replace('newline', '\n').replace('\n ', '\n'))

Answer 1

At a high level, I think you're overusing regexes. You can split stuff into words earlier and then clean them, and you don't need to treat newlines as an explicit token. If you do want to split them as an explicit token called newline to protect it from future regexes, you can handle that when you're constructing the markov chain by having a dedicated class, or just a string \n.

class NewlineToken:
    pass

I would take your input text and split it into words first, but do it a line at a time so I don't need to read the whole file into memory. Also, don't use file as the name of a temporary file since it shadows a builtin.

The split method when called with no arguments splits on whitespace.

with open("corpus.txt") as fh:
    for line in fh:
        words = line.split()
        # do stuff with words

After that you can strip punctuation from your words and lowercase it.

word.lower().translate(None, string.punctuation)

I'd probably wrap that all up in a generator that looks something like this.

def stripped_words(path):
    with open(path) as fh:
        for line in fh:
            words = line.split()
            for word in words:
                 yield word.lower().translate(
                         None, 
                         string.punctuation
                 )

More about string.punctuation here https://stackoverflow.com/questions/265960/best-way-to-strip-punctuation-from-a-string-in-python

The markov chain construction is good. zip automatically truncates the longer iterator and .setdefault is really useful. I actually didn't know about that until reading this. You'd have to modify it a little bit if you're iterating over words from your file rather than working with an array.

This line selects a word uniformly at random regardless of how many times it appears in the text rather than selecting a word with probability equal to its frequency. I'm not sure what the requirements are, but it is something that jumped out at me.

# more frequent words aren't more likely to be chosen here
current_word = random.choice(list(markov_chain.keys()))

Some miscellaneous comments:

You don't need to escape everything in your first character class.

text = re.sub(r'([,.!?;:])',r' \1 ', text)

In this line

text = re.sub(r'[^a-zA-Z0-9|\,|\.|\!|\?|\;|\:]',r' ',text)

| is part of your character class, I'm not sure whether you intended it because it isn't escaped like your other non-alphanumeric characters are.

Answer 2

Use functions

As it is, your code performs 3 tasks: getting a list of words out of a text, building the chain out of this words, and testing the chain by building a sentence; but it is rather difficult to get it at a glance as the whole code is at the top-level of the file. You should split your code into functions to ease reusability and testing. You may also use the if __name__ == '__main__': construct to separate the core of the code from the tests.

Use collections

setdefault is great for the task at hand but it forces you to build the default object each time, slowing down things a bit. Appending to a list in a loop is also deemed innefficient.

Instead, you could start by using a collections.defaultdict rather than setdefault, letting Python optimize the building of the empty list only for missing keys. Second, I would favor the usage of collections.Counter instead of a list: inserting existing elements would be improved.

Of course switching from list to Counter would mean changing the testing code. I think it's a good idea to provide either a function, or subclassing Counter to allow for easily retrieval of a word depending on its probability.

Use itertools

As slicing a list copy the slice in memory, it can be costy. So you need something more efficient for your loop. And if you're looking for efficient iterators, you should look into itertools. One of the recipe, pairwize, is exactly what you need here.

Proposed improvements

import random
import re
import collections
import itertools


class MarkovCounter(collections.Counter):
    @property
    def total(self):
        return sum(self.values())

    def probability(self, key):
        return self[key] / self.total

    def probabilities(self):
        total = self.total
        for key, value in self.items():
            yield key, value / total

    def pick_most_common(self):
        (most_common, _), = self.most_common(1)
        return most_common

    def pick_at_random(self):
        n = random.randint(0, self.total)
        return next(itertools.islice(self.elements(), n, None))


def pairwize(iterable):
    first, second = itertools.tee(iterable)
    next(second, None)
    return zip(first, second)


def get_words(filename):
    with open(filename) as file:
        text = file.read().replace('\n', ' newline ')
    
    text = re.sub(r'([\,\.\!\?\;\:])',r' \1 ',text)
    text = re.sub(r'[^a-zA-Z0-9|\,|\.|\!|\?|\;|\:]',r' ',text)
    return [word.lower().strip() for word in text.split(' ') if word.strip() != '']
    

def create_chain(words):
    markov_chain = collections.defaultdict(MarkovCounter)
    for current_word, next_word in pairwize(words):
        markov_chain[current_word][next_word] += 1
    return markov_chain
    

if __name__ == '__main__':
    text = get_words("corpus.txt")
    markov_chain = create_chain(text)
    current_word = random.choice(text)

    sentence = []
    for _ in range(15):
        sentence.append(current_word)
        if current_word in markov_chain:
            current_word = markov_chain[current_word].pick_at_random()
        else:
            break

    print(' '.join(sentence).replace('newline', '\n').replace('\n ', '\n'))

I also dropped the equivalent of if len(markov_chain[current_word]) > 0 because it is always True when current_word is in markov_chain.