markovbot/markov/chain.go

// Package markov implements an order-N Markov chain over word sequences,
// suitable for training on a corpus and generating text from it.
//
// A chain of order N keys each transition on the previous N tokens. Higher
// orders reproduce the source more faithfully and diverge less; order 1 is
// maximally inventive and usually incoherent. Order 2 is the usual
// compromise and is the default.
package markov

import (
	"errors"
	"fmt"
	"math/rand"
	"strings"
	"sync"
)

// start and end are sentinels marking the boundaries of a training sample.
// They cannot collide with real tokens because tokenisation never produces
// a token containing a space.
const (
	start = "\x00start"
	end   = "\x00end"
)

// ErrUntrained is returned when generation is attempted on a chain that has
// seen no input.
var ErrUntrained = errors.New("markov: chain has not been trained")

// Chain is an order-N Markov chain. It is safe for concurrent use.
type Chain struct {
	order int

	mu sync.RWMutex
	// freq maps a prefix of `order` tokens to the observed successors and
	// their counts. Counts, not a flat list, so repeated training does not
	// grow memory without bound and sampling stays proportional to
	// observed frequency.
	freq map[string]map[string]int
	// total caches the sum of counts per prefix so sampling avoids a
	// second pass over the successor map.
	total map[string]int
}

// New returns an empty chain of the given order. Order must be at least 1.
func New(order int) (*Chain, error) {
	if order < 1 {
		return nil, fmt.Errorf("markov: order must be >= 1, got %d", order)
	}
	return &Chain{
		order: order,
		freq:  make(map[string]map[string]int),
		total: make(map[string]int),
	}, nil
}

// Order returns the chain's order.
func (c *Chain) Order() int { return c.order }

// Train adds a single sample (a sentence, message or line) to the chain.
// Samples shorter than the order still contribute: the prefix is padded
// with start sentinels, so even a one-word sample teaches an opening.
func (c *Chain) Train(text string) {
	tokens := Tokenize(text)
	if len(tokens) == 0 {
		return
	}

	// Pad so that generation has a well-defined entry point and so the
	// final token is known to be terminal.
	padded := make([]string, 0, c.order+len(tokens)+1)
	for range c.order {
		padded = append(padded, start)
	}
	padded = append(padded, tokens...)
	padded = append(padded, end)

	c.mu.Lock()
	defer c.mu.Unlock()
	for i := 0; i+c.order < len(padded); i++ {
		prefix := strings.Join(padded[i:i+c.order], " ")
		next := padded[i+c.order]
		if c.freq[prefix] == nil {
			c.freq[prefix] = make(map[string]int, 4)
		}
		c.freq[prefix][next]++
		c.total[prefix]++
	}
}

// TrainAll adds every sample in texts.
func (c *Chain) TrainAll(texts []string) {
	for _, t := range texts {
		c.Train(t)
	}
}

// Generate produces text of at most maxWords tokens, using rng for sampling.
//
// Passing a seeded *rand.Rand makes output deterministic, which is what
// makes this testable; callers wanting variety should seed from the clock.
func (c *Chain) Generate(rng *rand.Rand, maxWords int) (string, error) {
	return c.GenerateFrom(rng, "", maxWords)
}

// GenerateFrom produces text seeded with the given prefix text. When the
// seed is unknown to the chain, generation falls back to a fresh start
// rather than failing: a bot that answers something is more useful than one
// that refuses.
func (c *Chain) GenerateFrom(rng *rand.Rand, seed string, maxWords int) (string, error) {
	if rng == nil {
		return "", errors.New("markov: nil rng")
	}
	if maxWords <= 0 {
		return "", fmt.Errorf("markov: maxWords must be > 0, got %d", maxWords)
	}

	c.mu.RLock()
	defer c.mu.RUnlock()
	if len(c.freq) == 0 {
		return "", ErrUntrained
	}

	window := c.seedWindow(seed)
	out := make([]string, 0, maxWords)
	for range maxWords {
		prefix := strings.Join(window, " ")
		next, ok := c.sample(rng, prefix)
		if !ok || next == end {
			break
		}
		out = append(out, next)
		window = append(window[1:], next)
	}
	return strings.Join(out, " "), nil
}

// seedWindow builds the initial prefix window from seed text, padding on
// the left with start sentinels when the seed is shorter than the order.
func (c *Chain) seedWindow(seed string) []string {
	window := make([]string, c.order)
	for i := range window {
		window[i] = start
	}
	tokens := Tokenize(seed)
	if len(tokens) == 0 {
		return window
	}
	if len(tokens) > c.order {
		tokens = tokens[len(tokens)-c.order:]
	}
	copy(window[c.order-len(tokens):], tokens)

	// An unseen seed window would dead-end immediately; fall back to a
	// fresh start so the caller still gets output.
	if _, ok := c.freq[strings.Join(window, " ")]; !ok {
		for i := range window {
			window[i] = start
		}
	}
	return window
}

// sample picks a successor for prefix with probability proportional to how
// often it was observed. Caller must hold at least a read lock.
func (c *Chain) sample(rng *rand.Rand, prefix string) (string, bool) {
	successors, ok := c.freq[prefix]
	if !ok || len(successors) == 0 {
		return "", false
	}
	total := c.total[prefix]
	if total <= 0 {
		return "", false
	}

	// Iteration order over a Go map is randomised, which would make output
	// non-reproducible even with a seeded rng. Walking a sorted key list
	// keeps sampling a pure function of the rng.
	n := rng.Intn(total)
	for _, k := range sortedKeys(successors) {
		n -= successors[k]
		if n < 0 {
			return k, true
		}
	}
	return "", false
}

// Prefixes returns the number of distinct prefixes learned.
func (c *Chain) Prefixes() int {
	c.mu.RLock()
	defer c.mu.RUnlock()
	return len(c.freq)
}

// Tokenize splits text into tokens on whitespace. Sentinel-lookalike tokens
// are dropped so a hostile corpus cannot inject chain boundaries.
func Tokenize(text string) []string {
	fields := strings.Fields(text)
	out := fields[:0]
	for _, f := range fields {
		if f == start || f == end {
			continue
		}
		out = append(out, f)
	}
	return out
}