Task: Markov chain text generator




Go: package main

import (
	"bufio"
	"flag"
	"fmt"
	"io"
	"log"
	"math/rand"
	"os"
	"strings"
	"time"
	"unicode"
	"unicode/utf8"
)

func main() {
	log.SetFlags(0)
	log.SetPrefix("markov: ")
	input := flag.String("in", "alice_oz.txt", "input file")
	n := flag.Int("n", 2, "number of words to use as prefix")
	runs := flag.Int("runs", 1, "number of runs to generate")
	wordsPerRun := flag.Int("words", 300, "number of words per run")
	startOnCapital := flag.Bool("capital", false, "start output with a capitalized prefix")
	stopAtSentence := flag.Bool("sentence", false, "end output at a sentence ending punctuation mark (after n words)")
	flag.Parse()

	rand.Seed(time.Now().UnixNano())

	m, err := NewMarkovFromFile(*input, *n)
	if err != nil {
		log.Fatal(err)
	}

	for i := 0; i < *runs; i++ {
		err = m.Output(os.Stdout, *wordsPerRun, *startOnCapital, *stopAtSentence)
		if err != nil {
			log.Fatal(err)
		}
		fmt.Println()
	}
}
















type Markov struct {
	n           int
	capitalized int 
	suffix      map[string][]string
}



func NewMarkovFromFile(filename string, n int) (*Markov, error) {
	f, err := os.Open(filename)
	if err != nil {
		return nil, err
	}
	defer f.Close() 
	return NewMarkov(f, n)
}



func NewMarkov(r io.Reader, n int) (*Markov, error) {
	m := &Markov{
		n:      n,
		suffix: make(map[string][]string),
	}
	sc := bufio.NewScanner(r)
	sc.Split(bufio.ScanWords)
	window := make([]string, 0, n)
	for sc.Scan() {
		word := sc.Text()
		if len(window) > 0 {
			prefix := strings.Join(window, " ")
			m.suffix[prefix] = append(m.suffix[prefix], word)
			
			if isCapitalized(prefix) {
				m.capitalized++
			}
		}
		window = appendMax(n, window, word)
	}
	if err := sc.Err(); err != nil {
		return nil, err
	}
	return m, nil
}





func (m *Markov) Output(w io.Writer, n int, startCapital, stopSentence bool) error {
	
	
	
	bw := bufio.NewWriter(w)

	var i int
	if startCapital {
		i = rand.Intn(m.capitalized)
	} else {
		i = rand.Intn(len(m.suffix))
	}
	var prefix string
	for prefix = range m.suffix {
		if startCapital && !isCapitalized(prefix) {
			continue
		}
		if i == 0 {
			break
		}
		i--
	}

	bw.WriteString(prefix) 
	prefixWords := strings.Fields(prefix)
	n -= len(prefixWords)

	for {
		suffixChoices := m.suffix[prefix]
		if len(suffixChoices) == 0 {
			break
		}
		i = rand.Intn(len(suffixChoices))
		suffix := suffixChoices[i]
		
		bw.WriteByte(' ') 
		if _, err := bw.WriteString(suffix); err != nil {
			break
		}
		n--
		if n < 0 && (!stopSentence || isSentenceEnd(suffix)) {
			break
		}

		prefixWords = appendMax(m.n, prefixWords, suffix)
		prefix = strings.Join(prefixWords, " ")
	}
	return bw.Flush()
}

func isCapitalized(s string) bool {
	
	
	
	r, _ := utf8.DecodeRuneInString(s)
	return unicode.IsUpper(r)
}

func isSentenceEnd(s string) bool {
	r, _ := utf8.DecodeLastRuneInString(s)
	
	
	
	return r == '.' || r == '?' || r == '!'
}

func appendMax(max int, slice []string, value string) []string {
	
	
	
	
	
	
	
	
	
	
	if len(slice)+1 > max {
		n := copy(slice, slice[1:])
		slice = slice[:n]
	}
	return append(slice, value)
}




Python: import random, sys

def makerule(data, context):
    
    rule = {}
    words = data.split(' ')
    index = context
    
    for word in words[index:]:
        key = ' '.join(words[index-context:index])
        if key in rule:
            rule[key].append(word)
        else:
            rule[key] = [word]
        index += 1

    return rule


def makestring(rule, length):    
    
    oldwords = random.choice(list(rule.keys())).split(' ') 
    string = ' '.join(oldwords) + ' '
    
    for i in range(length):
        try:
            key = ' '.join(oldwords)
            newword = random.choice(rule[key])
            string += newword + ' '

            for word in range(len(oldwords)):
                oldwords[word] = oldwords[(word + 1) % len(oldwords)]
            oldwords[-1] = newword

        except KeyError:
            return string
    return string


if __name__ == '__main__':
    with open(sys.argv[1], encoding='utf8') as f:
        data = f.read()
    rule = makerule(data, int(sys.argv[2]))
    string = makestring(rule, int(sys.argv[3]))
    print(string)