Below is the file 'mk2.py' from this revision. You can also download the file.

#!/usr/bin/env python

import cPickle
import random
random.seed()
import heapq
import math
import sys

class MarkovState(object):
    def __init__(self, state):
        self.state = state
        self.h = None
        self.total = 0
        self.scores = {}

    def increment(self, token):
        self.total += 1
        self.scores[token] = self.scores.get(token, 0) + 1
        self.h = None

    def __entropy(self):
        return -1 * sum(map(lambda p: p * math.log(p, 2),
                        map(lambda x: (self.scores[x] / float(self.total)), self.scores)))

    def entropy(self):
        if self.h == None:
            self.h = self.__entropy()
        return self.h

    def __repr__(self):
        return repr(self.scores)

    def __cmp__(self, other):
        if other == None:
            return -1
        return cmp(other.entropy(), self.entropy())

class MarkovChain(object):
    def __init__(self, length):
        self.length = length
        self.clear()

    def update(self, gen):
        buffer = []
        for token in gen:
            self.stash.append(token)
            if len(buffer) == self.length:
                tbuffer = tuple(buffer)
                if self.states.has_key(tbuffer):
                    state = self.states[tbuffer]
                else:
                    state = self.states[tbuffer] = MarkovState(tbuffer)
                state.increment(token)
                buffer = buffer[1:]
            buffer.append(token)

    def clear(self):
        self.states = {}
        self.stash = []

    def random_next(self, from_state):
        def next_state(token):
            return from_state.state[:-1] + (token,)
        # eliminate dead-ends
        def not_dead_end(token):
            return self.states.has_key (next_state (token))
        possible = filter (not_dead_end, from_state.scores.keys())
#        print >>sys.stderr, (from_state, possible)
        if not possible:
            return None
        total = sum (map (lambda s: from_state.scores[s], possible))
        choice = random.randrange(0, total)
        for k in possible:
            total -= from_state.scores[k]
            if total <= 0:
                return self.states[next_state(k)]
        raise Exception("Unreachable")

    def upchunk(self):
        while True:
            to_upchunk = self.__select_upchunk()
            if to_upchunk == None:
                break
            stash_copy = self.stash
            self.clear()
            self.update(self.__upchunk_gen (stash_copy, to_upchunk))
            del stash_copy

    def __select_upchunk(self):
        q = []
        keys = self.states.keys()
        if len(keys) == 0:
            return None
        for idx, tokens in enumerate(keys):
            state = self.states[tokens]
            heapq.heappush(q, state)
        cutoff = math.log (len (keys), 2) / 4
        candidate = heapq.heappop(q)
        print >>sys.stderr, "best entropy vs. cutoff is: %s :: %.2f vs. cutoff %.2f" % (candidate.state, candidate.entropy(), cutoff)
        if candidate.entropy() < cutoff:
            return None
        else:
            return candidate.state

    def __upchunk_gen(self, gen, to_upchunk):
        buffer = []
        for i in gen:
            buffer.append(i)
            if len(buffer) == len(to_upchunk):
                if tuple(buffer) == to_upchunk:
                    buffer = [ ''.join(buffer) ]
                else:
                    to_yield, buffer = buffer[0], buffer[1:]
                    yield to_yield
        for i in buffer:
            yield i

    def pprint(self):
        from pprint import pprint
        pprint(chain.states)

def simple_gen(fname):
    for line in open(fname, 'rb'):
        for char in line:
            yield char
#        for word in line.split():
#            yield word.lower()

if __name__ == '__main__':
    chain = MarkovChain(2)
    for infile in sys.argv[1:]:
        print >> sys.stderr, "Reading input file:", infile
        chain.update(simple_gen (infile))
    chain.upchunk()
    print >>sys.stderr, "processing produced", len(chain.states.keys()), "states."
    cPickle.dump(chain, sys.stdout, protocol=2)