compose_template.py

import os
import re
import string
import random

from graph import Graph, Vertex

def get_words_from_text(text_path):
    with open(text_path, "rb") as file:
        text = file.read().decode("utf-8")
        text = " ".join(text.split())
        text = text.lower()
        text = text.translate(str.maketrans("", "", string.punctuation))
    words = text.split()
    return words[:1000]

def make_graph(words):
    g = Graph()
    prev_word = None
    for word in words:
        word_vertex = g.get_vertex(word)
        if prev_word:
            prev_word.increment_edge(word_vertex)
        prev_word = word_vertex
    g.generate_probability_mappings()
    return g

def compose(g, words, length=50):
    composition = []
    current = g.get_vertex(random.choice(words))
    for _ in range(length):
        composition.append(current.value)
        current = g.get_next_word(current)
    return composition

def main():
    words = get_words_from_text("hp_socerer_stone.txt")
    graph = make_graph(words)
    composition = compose(graph, words, 100)
    print(" ".join(composition))


if __name__ == '__main__':
    main()