Генерация текста с помощью цепи Маркова на F#

16.09.2014 at 17:59


Продолжаю изучать F#. В прошлый раз я написал парсер хабра, теперь на основе него буду писать генератор текста.

Алгоритм генерации

На входе в алгоритм — текст, в нашем случае — спарсенные с хабра статьи. Весь текст разбивается на последовательности из трех слов. Каждое новое слово в генерируемом тексте подбирается на основе двух предыдущих следующим образом: берутся все тройки, которые начинаются с этих слов и из них выбирается случайная. Вероятность выбора каждой тройки пропорциональна частоте ее использования в исходном тексте. После того, как мы выбрали тройку, из нее берется последнее слово и добавляется в генерируемый текст.

Код генератора текста

Сначала импортируем нужные пространства имен:

open System
open System.IO

Определяем тип данных, в котором будут храниться тройки слов и их частота:

type CountedListElement = {Words: string list; Count: int}

Составляем список файлов с текстом:

let fileNames = Directory.GetFiles(@"D:\temp\habr") |> Array.toList 

Определяем функцию для получения последнего слова из тройки:

let lastWord (x: CountedListElement) = 
  x.Words.Tail.Tail.Head

Парсим файл с текстом:

let rec parseFile (fileContent: string list) = 
  match fileContent with
  | x :: y :: z :: tail -> [x; y; z] :: (parseFile fileContent.Tail) 
  | _ -> []

Эта функция принимает на вход список слов, а возвращает — список троек слов.

Читаем содержимое файла и разделяем его на слова:

let readFile fileName = 
  let fileContent = File.ReadAllText(fileName).Replace("...", ".").Replace(".", " . ").Replace("!", " ! ").Replace("?", " ? ").Replace(Environment.NewLine, " ")
  "" :: "" :: (fileContent.Split([|" "|], StringSplitOptions.RemoveEmptyEntries) |> Array.toList)

Все знаки препинания обрабатываются как отдельные слова. В начало списка слов добавляются 2 пустые строки. Так мы отмечаем начало текста.

Объединяем 2 предыдущие функции:

let prepareFile fileName =
  let fileContent = readFile fileName 
  parseFile fileContent

Добавляем тройку слов в список троек:

let addElementToFrequencyList (fList: CountedListElement list) (element: string list) = 
  match fList with
  | head :: tail when head.Words = element -> {Words = head.Words; Count = head.Count + 1} :: tail 
  | _ -> { Words = element; Count = 1 } :: fList

Если тройка в списке уже есть — увеличиваем частоту на 1. Если нет — добавляем.

Служебная функция для сравнения списков строк:

let rec compareStringLists (x: string list) (y: string list) = 
  match x, y with
  | q::tailx, w::taily when q = w -> compareStringLists tailx taily
  | q::tailx, w::taily -> Operators.compare q w
  | _, _ -> 0

Создаем список троек с частотами:

let generationBase = 
  fileNames
  |> List.map (fun x -> prepareFile x) 
  |> List.concat
  |> List.sortWith  (fun x y -> compareStringLists x y)
  |> List.fold (fun acc x -> addElementToFrequencyList acc x) ([]:CountedListElement list)

Создаем сгенерированный текст:

let initialText = [""; ""]

Пока в нем только 2 пустые строки — так мы обозначали начало текста при парсинге.

Выбираем из списка троек случайный элемент:

let rec selectCountedListElementByFrequency (lst: CountedListElement list) (frequency: int) = 
  if lst.Head.Count > frequency
  then
    lst.Head
  else
    selectCountedListElementByFrequency lst.Tail (frequency - lst.Head.Count)

Функция, которая генерирует следующее слово:

let findNextWord word1 word2 = 
  let usableWords = generationBase |> List.filter (fun x -> x.Words.Head = word1 && x.Words.Tail.Head = word2)
  if usableWords.Length > 0
  then
    let totalFrequency = usableWords |> List.fold (fun acc x -> acc + x.Count) 0
    let selectedFrequency = (new Random()).Next(totalFrequency) 
    lastWord (selectCountedListElementByFrequency usableWords selectedFrequency)
  else
    ""

Функция, которая добавляет к тексту следующее слово:

let rec addWords (text: string list) (count: int) = 
  match text, count with
  | _, count when count <= 0 -> text
  | x, _ when (x.Length < 2) -> text
  | text, count -> addWords ((findNextWord text.Tail.Head text.Head)::text) (count - 1)

Новое слово добавляется в начало списка, так как списки в F# хранят указатель на начало списка, но не хранят указатель на конец списка.

Генерируем 500 слов и сохраняем в файл:

let resultList = addWords initialText 500 |> List.rev 
let resultText = String.Join(" ", resultList)
File.AppendAllText(@"D:\result.txt", resultText)