markov_chain.ex

defmodule MarkovChain do
  def create_chain(source, ngram_length) do
    source
    |> Stream.chunk_every(ngram_length + 1, 1)
    |> Stream.map(fn chunk -> Enum.split(chunk, ngram_length) end)
    |> Stream.scan(%{}, fn {ngram, word}, acc ->
      Map.update(acc, ngram, word, fn val -> val ++ word end)
    end)
  end

  def next_element(chain, preceding) do
    chain
    |> Map.get(preceding, [])
    |> Enum.random()
  end
end

markov_chain_test.exs

defmodule MarkovChainTest do
  use ExUnit.Case
  doctest MarkovChain

  describe "create_chain" do
    test "returns correct map for names" do
      words = "Tom Bombadil" |> String.graphemes()

      assert %{
               ["a"] => ["d"],
               [" "] => ["B"],
               ["B"] => ["o"],
               ["T"] => ["o"],
               ["b"] => ["a"],
               ["d"] => ["i"],
               ["i"] => ["l"],
               ["l"] => [],
               ["m"] => [" ", "b"],
               ["o"] => ["m", "m"]
             } == MarkovChain.create_chain(words, 1) |> Enum.to_list() |> List.last()
    end

    test "returns correct map for words" do
      words = "a long time ago in a galaxy far far away" |> String.split()

      assert %{
               ["a"] => ["long", "galaxy"],
               ["long"] => ["time"],
               ["time"] => ["ago"],
               ["ago"] => ["in"],
               ["in"] => ["a"],
               ["galaxy"] => ["far"],
               ["far"] => ["far", "away"],
               ["away"] => []
             } == MarkovChain.create_chain(words, 1) |> Enum.to_list() |> List.last()
    end
  end
end