markov-chain-model.py

# coding: utf-8

# Hans Kamin
# Spring 2017

# Bigram Markov Chain Model

def train_markov_chain(lyrics):
    """
    Args:
      - lyrics: a list of strings, where each string represents
                the lyrics of one song by an artist.

    Returns:
      A dict that maps a tuple of 2 words ("bigram") to a list of
      words that follow that bigram, representing the Markov
      chain trained on the lyrics.
    """

    # Initialize the beginning of our chain.
    chain = {
        (None, "<START>"): []
    }

    for lyric in lyrics:
        # Replace newline characters with our tag.
        lyric_newlines = lyric.replace('\n', ' <N> ')
        # Create a tuple representing the most recent (current) bigram.
        last_2 = (None, "<START>")
        for word in lyric_newlines.split():
            # Add the word as one that follows the current bigram.
            chain[last_2].append(word)
            # Shift the current bigram to account for the newly added word.
            last_2 = (last_2[1], word)
            if last_2 not in chain:
                chain[last_2] = []
        chain[last_2].append("<END>")

    return chain

# Load the pickled lyrics object that we created earlier.
import pickle
lyrics = pickle.load(open("lyrics.pkl", "rb"))

# Train a Markov Chain over all of Logic's lyrics.
chain = train_markov_chain(lyrics)


import random

def generate_new_lyrics(chain):
    """
    Args:
      - chain: a dict representing the Markov chain,
               such as one generated by generate_new_lyrics()

    Returns:
      A string representing the randomly generated song.
    """

    # a list for storing the generated words
    words = []
    # generate the first word
    word = random.choice(chain[(None, "<START>")])
    words.append(word)

    # Begin with the first bigram in our chain.
    last_2 = (None, "<START>")
    while words[-1] != "<END>":
        # Generate the next word.
        word = random.choice(chain[last_2])
        words.append(word)
        # Shift the current bigram to account for the newly added word.
        last_2 = (last_2[1], words[-1])

    # Join the words together into a string with line breaks.
    lyrics = " ".join(words[:-1])
    return "\n".join(lyrics.split("<N>"))



print(generate_new_lyrics(chain))