autocomplete.py

import random

# Change these files to the ones you want.
SOURCE_FILE = "origs/therepublic.txt"
OUT_FILE = "out/therepublic.txt"


def loadFile(fn):
    with open(fn, encoding="ISO-8859-1") as f:
        return f.read()


def splitFile(text):
    return [word for word in text.split(" ") if len(word) > 0]


class Chain:
    def __init__(self):
        self.wordList = {}

    def addWordsToList(self, prev, word):
        if prev in self.wordList.keys():
            if word in self.wordList[prev].keys():
                self.wordList[prev][word] += 1
            else:
                self.wordList[prev][word] = 1

        else:
            self.wordList[prev] = {word: 1}

    def getRandomChoiceList(self, word):
        potentialWords = self.wordList[word]
        weightedChoices = []
        for word in potentialWords.keys():
            wordFreq = potentialWords[word]
            for i in range(wordFreq ** 2):
                weightedChoices.append(word)
        return weightedChoices

    def nextWord(self, currentWord):
        return random.choice(self.getRandomChoiceList(currentWord))

    def generateChain(self, startWord, length):
        chain = []
        word = startWord
        for i in range(length):
            chain.append(word)
            word = self.nextWord(word)
        return chain


def fillChain(chain, wordList):
    for i in range(len(wordList)-1):
        prevWord = wordList[i]
        word = wordList[i+1]

        chain.addWordsToList(prevWord, word)


def saveToFile(text, fn):
    with open(fn, "w+") as f:
        f.write(text)


oneWordChain = Chain()

print("Loading text...")
text = loadFile(SOURCE_FILE)

print("Parsing text...")
splitText = splitFile(text)

print("Filling chain...")
fillChain(oneWordChain, splitText)

print("Generating...")
saveToFile(" ".join(oneWordChain.generateChain(
    'the', 1000)).replace("\n", " "), OUT_FILE)