Markov-chain sentence generator in Python

Question

I wrote a Markov-chain based sentence generator as my first non-trivial Python program. I mainly used C before, so I probably have ignored a lot of Python conventions and features, so any advice would be appreciated.

text-gen.py

import sys
import random

class MarkovChain:
    # Class constant that serves as an initial state for the Markov chain
    START = ""

    # The Markov chain is modelled as a directed graph,
    # with the START state acting as the only source,
    # and the tranisition probabilities as the graph weights.
    #
    # The graph is implemented using an adjacency list,
    # which in turn is implemented as a dictionary of dictionaries.
    #
    # self.adjList is a dictionary keyed by all words of the text
    # or START (the states). For each key/state, it contains
    # another dictionary indexed by the words of the text
    # that succeed the key in the text (the next states in the chain),
    # and for each of those words/next states the dictionary contains
    # the transition probability from the present state to them.
    #
    # This implementation was chosen because of it being easy to code,
    # and offering an easy way to iterate on both the probabilities and
    # the words/next states of each dictionary using items().
    def __init__(self, file):
        self.adjList = {}

        # The totals dictionary is used in calculating the probabilities,
        # for every word in the text/chain state it contains the total
        # number of transitions from it to another state.
        totals = {}


        # Start by insering the initial state to the structures
        self.adjList[MarkovChain.START] = {}
        totals[MarkovChain.START] = 0

        # prev: Contains the previously encountered word or the START state,
        # initialized to the START state.
        prev = MarkovChain.START


        for line in file:
            for word in line.split():
                # If the word ends with a terminating punctuation mark,
                # ignore the mark, and treat the word as a terminating state as
                # it does not preceed another word in the current sentence.
                # So prev is set to START, in order for the text model
                # to account for the fact that some words start sentences
                # more frequently than others (not all words are next states of START).
                endsTerm = word[-1] == "." or word[-1] == "?" or word[-1] == "!"
                if (endsTerm):
                    word = word[0:-1]

                # If this is the first time the word is encountered,
                # add it to the adjacency list, and initialize its dictionary
                # and transition total.
                if (word not in self.adjList):
                    self.adjList[word] = {}
                    totals[word] = 0

                # If this is the first time the prev->word transition
                # was detected, initialize the prev->word transition frequency to 1,
                # else increment it.
                if (word in self.adjList[prev]):
                    self.adjList[prev][word] += 1 
                else:
                    self.adjList[prev][word] = 1

                # There is a prev->word state transition, so increment
                # the total transition number of the prev state.
                totals[prev] += 1

                if (endsTerm):
                    prev = START

        # Using total, convert the transition frequencies
        # to transition probabilities.
        for word, neighbors in self.adjList.items():
            for name in neighbors:
                neighbors[name] /= totals[word]


    # chooseNextWord: Chooses the next state/word,
    # by sampling the non uniform transition probability distribution
    # of the current word/state. 
    def chooseNextWord(self, curWord):
        # Convert the dict_keys object to a list
        # to use indexing
        nextWords = list(self.adjList[curWord].keys())

        # Sampling is done through linear search.
        for word in nextWords[0:-1]:
            prob = self.adjList[curWord][word]
            roll = random.random()
            if (roll <= prob):
                return word

        # If none of the first N-1 words were chosen,
        # only the last one was left.
        return nextWords[-1]


    # genSentence: Generates a sentence. If a positive
    # limit is not provided by the caller, the sentences grow to
    # an arbitrary number of words, until the last word of a sentence/a terminal state
    # is reached.
    def genSentence(self, limit = 0):
        sentence = ""

        curWord = self.chooseNextWord(MarkovChain.START)
        sentence += curWord + " "

        if (limit > 0):
            wordsUsed = 1
            while (wordsUsed < limit and self.adjList[curWord]):
                curWord = self.chooseNextWord(curWord)
                sentence += curWord + " "
                wordsUsed += 1
        else:
            while (self.adjList[curWord]):
                curWord = self.chooseNextWord(curWord)
                sentence += curWord + " "

        return sentence



if (__name__ == "__main__"):
    if (len(sys.argv) < 3):
        print("Not enough arguements, run with python3 text-gen.py <input-filename> <sentence-num>")
        sys.exit(1)

    try:
        with open(sys.argv[1], "r") as f:
            markov = MarkovChain(f)

    except OSError as error:
        print(error.strerror)
        sys.exit(1)

    # Generate and print as many sentences as asked.
    for k in range(0, int(sys.argv[2])):
        print(markov.genSentence(20) + "\n")

Answer 1

• chooseNextWord distorts the probabilities.

  For example, consider a list of 3 words with the inherent probabilities \$\frac{1}{3}\$, \$\frac{1}{3}\$, \$\frac{1}{3}\$. The first word is selected with the probability \$\frac{1}{3}\$. The second, however is selected with probability \$\frac{2}{9}\$ (\$\frac{2}{3}\$ that the first word was not selected, times \$\frac{1}{3}\$ that it is selected in the second round). The third one has \$1 - \frac{1}{3} - \frac{2}{9} = \frac{4}{9}\$ chance.

  A standard approach is to compute an accumulated sums of probabilities (in the constructor), then to choose a word roll once, and search for a value just above the rolled one.

• The code may benefit from using defaultdict rather than a plain dictionaries. Lesser ifs is better.

• Nitpicking. You may want to account for possible typos, such as a space between a word and a terminating punctuation.

Answer 2

endTerms

When a word ends with an endTerm, think you need to include an START or END symbol in adjList. Most words can appear anywhere in a sentence. So it is unlikely that you can end a sentence only when words don't have any follow-on words. Include the START/END symbol in adjList and the Markov process can also end a sentence.

the standard library

collections.defaultdict provides a dictionary that when you attempt to access a new key automatically initializes the new key to a default value.

collections.Counter provides a dictionary that counts things.

random.choices selects items from a population according to specified weights.

import collections
import random

class MarkovChain:
    START = ""

    def __init__(self, file):
        adjList = collections.defaultdict(collections.Counter)

        # this inserts START into the defaultdict
        adjList[MarkovChain.START]

        prev = MarkovChain.START

        for line in file:
            for word in line.split():
                endsTerm = word[-1] in ('.', '?', '!')

                if (endsTerm):
                    word = word[:-1]

                adjList[prev].update([word])

                if endsTerm:
                    # mark the end of a sentence
                    adjList[word].update([MarkovChain.START])
                    prev = MarkovChain.START
                else:
                    prev = word

        # convert defaultdict to a regular dict
        # the values are a tuple: ([follow words], [counts])
        # for use in random.choices() in chooseNextWord()
        self.adjList = {k:(list(v.keys()), list(v.values()))
                        for k,v in adjList.items()}

        #print(self.adjList)


    def chooseNextWord(self, word):
        # random.choices returns a list, hence the [0]
        return random.choices(*self.adjList[word])[0]


    def genSentence(self, limit = 0):
        sentence = []
        curWord = MarkovChain.START

        while True:
            curWord = self.chooseNextWord(curWord)
            sentence.append(curWord)

            if 0 < limit < len(sentence) or curWord == MarkovChain.START:
                break

        return ' '.join(sentence).strip()

Answer 3

Style

• The long comment at the beginning of the class is good class documentation and should be made into a docstring. The same applies for functions.
• Avoid variables in all caps, like START (unless they are constants).
• Functions and variables should generally use snake-case, if you want to follow Python conventions as in PEP 8.

Answer 4

Let's look at this one line, which shows quite a few areas where your code could be improved:

if (word in self.adjList[prev]):

• As noted elsewhere, Python has an official style-guide, PEP8, which recommends using lower_case for variables and functions.
• Basically all Python keywords (at least if, elif, else, for, while, in, with, except, return, yield, yield from, print in Python 2) take an expression as argument. Expressions do not need to be surrounded by parentheses.
• Whenever you do if X in Y you should know what the time complexity of that statement is. For tuple, str and list it is just a linear scan, so it is linear. For set and dict it is (amortized) constant due to the use of hashes.
• adjList is not actually a list, but a dict! This is why Hungarian notation (putting the type in the name) is discouraged and in addition duck typing is encouraged. Here you could just call it adjacancies.
• Take an afternoon and work through the Python standard library. I would recommend at least collections (where you can find defaultdict and Counter as recommend Ed in other answers, itertools (and more_itertools for bonus points, although it is not in the standard library), str, functools, math, random and pathlib as a start.

Answer 5

For long blocks of informational text at the beginning of a class or method definition, you should use docstrings instead of comments, as per PEP 8. This makes your descriptions automatically available from the help() function.

Example:

def foo(bar):
    """Foo does something.
    This description can span multiple lines
    """
    return bar