Генератор текста на основе триграмм

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

Сухая теория

И так, наша задача сгенерировать текст. Это значит, нам нужно взять слова и выстроить их в определенном порядке. Как определить этот порядок? Мы можем пойти следующим образом: построить фразы, наиболее вероятные для русского языка. Но что значит вероятность фразы языка? С точки зрения здравого смысла это бред. Тем не менее, эту вероятность можно задать формально как вероятность возникновения последовательности слов в неком корпусе (наборе текстов). К примеру, вероятность фразы "счастье есть удовольствие без раскаяния" можно вычислить как произведение вероятностей каждого из слов этой фразы:

P = P(счастье) P(есть|счастье) P(удовольствие|счастье есть) P(без|счастье есть удовольствие) P(раскаяния|счастье есть удовольствие без)

Рассчитать вероятность P(счастье) дело нехитрое: нужно всего лишь посчитать сколько раз это слово встретилось в тексте и поделить это значение на общее число слов. Но рассчитать вероятность P(раскаяния|счастье есть удовольствие без) уже не так просто. К счастью, мы можем упростить эту задачу. Примем, что вероятность слова в тексте зависит только от предыдущего слова. Тогда наша формула для расчета фразы примет следующий вид:

P = P(счастье) P(есть|счастье) P(удовольствие|есть) P(без|удовольствие) P(раскаяния|без)

Уже проще. Рассчитать условную вероятность P(есть|счастье) несложно. Для этого считаем количество пар 'счастье есть' и делим на количество в тексте слова 'счастье':

P(есть|счастье) = C(счастье есть)/С(счастье)

В результате, если мы посчитаем все пары слов в некотором тексте, мы сможем вычислить вероятность произвольной фразы. А если мы сможем вычислить вероятность фразы, мы можем подобрать наиболее "правдоподобные" сочетания слов в данном тексте для автоматической генерации. Кстати, эти пары слов называются биграммы, а набор рассчитанных вероятностей - биграммной моделью.

Хочу отметить, что для нашего алгоритма мы используем не биграммы, а триграммы. Отличие в том, что условная вероятность слова определяется не по одному, а по двум предыдущим словам. То есть вместо P(удовольствие|счастье) мы будем вычислять P(удовольствие|счастье есть). Формула вычисления подобна формуле для биграмм:

P(удовольствие|счастье есть) = C(счастье есть удовольствие)/С(счастье есть)

Итак, генерацию текста можно провести следующим образом. Мы будем формировать отдельно каждое предложение, выполняя следующие шаги:

Практика

Для начала, нам нужно подготовить корпус, на котором мы будем тренировать свою модель. Я, к примеру, взял выборку из Льва Толстого на сайте lib.ru, и сформировал один текстовый файл (скачать можно здесь).

Из данного текста выделим необходимую нам последовательность слов (если данный код вызывает вопросы, прошу глянуть в статейку по обработке текста).

import re

r_alphabet = re.compile(u'[а-яА-Я0-9-]+|[.,:;?!]+')

def gen_lines(corpus):
    data = open(corpus)
    for line in data:
        yield line.decode('utf-8').lower()

def gen_tokens(lines):
    for line in lines:
        for token in r_alphabet.findall(line):
            yield token

lines = gen_lines('tolstoy.txt')
tokens = gen_tokens(lines)

Получившийся генератор tokens будет выдавать "очищенную" последовательность слов и знаков препинания. Однако, простая последовательность нам не интересна. Нам интересны тройки токенов (под токеном здесь понимается слово или знак препинания, т.е. некие атомарные элементы текста). Для этого добавим еще один генератор, на выходе которого будем иметь три подряд идущих токена.

import re

r_alphabet = re.compile(u'[а-яА-Я0-9-]+|[.,:;?!]+')

def gen_lines(corpus):
    data = open(corpus)
    for line in data:
        yield line.decode('utf-8').lower()

def gen_tokens(lines):
    for line in lines:
        for token in r_alphabet.findall(line):
            yield token

def gen_trigrams(tokens):
    t0, t1 = '$', '$'
    for t2 in tokens:
        yield t0, t1, t2
        if t2 in '.!?':
            yield t1, t2, '$'
            yield t2, '$','$'
            t0, t1 = '$', '$'
        else:
            t0, t1 = t1, t2

lines = gen_lines('tolstoy.txt')
tokens = gen_tokens(lines)
trigrams = gen_trigrams(tokens)

Метод gen_trigrams требует пояснения. Символы '$' используются для маркировки начала предложения. В дальнейшем, это делает проще выбор первого слова генерируемой фразы. В целом, метод действует следующим образом: он возвращает три подряд идущих токена, на каждой итерации сдвигаясь на один токен:

На входе: 
    'Счастье есть удовольствие без раскаяния'

На выходе: 
    итерация    токены
    0:          $ $ счастье
    1:          $ счастье есть
    2:          счастье есть удовольствие
    3:          есть удовольствие без
    ...

Далее, рассчитаем триграммную модель:

import re
from collections import defaultdict

r_alphabet = re.compile(u'[а-яА-Я0-9-]+|[.,:;?!]+')

...

def train(corpus):
    lines = gen_lines(corpus)
    tokens = gen_tokens(lines)
    trigrams = gen_trigrams(tokens)

    bi, tri = defaultdict(lambda: 0.0), defaultdict(lambda: 0.0)

    for t0, t1, t2 in trigrams:
        bi[t0, t1] += 1
        tri[t0, t1, t2] += 1

    model = {}
    for (t0, t1, t2), freq in tri.iteritems():
        if (t0, t1) in model:
            model[t0, t1].append((t2, freq/bi[t0, t1]))
        else:
            model[t0, t1] = [(t2, freq/bi[t0, t1])]
    return model

model = train('tolstoy.txt')

В первой части этого метода мы задаем генераторы. Далее, рассчитываем биграммы и триграммы (фактически, мы считаем количество одинаковых пар и троек слов в тексте). Далее, вычисляем вероятность слова в зависимости от двух предыдущих, помещая данное слово и его вероятность в словарь. Надо сказать, что это не самый оптимальный метод, так как идет значительный расход памяти. Но для небольших корпусов этого вполне достаточно.

Теперь у нас все готово для генерации текста. Следующая функция возвращает предложение.

...

model = train('tolstoy.txt')

def generate_sentence(model):
    phrase = ''
    t0, t1 = '$', '$'
    while 1:
        t0, t1 = t1, unirand(model[t0, t1])
        if t1 == '$': break
        if t1 in ('.!?,;:') or t0 == '$':
            phrase += t1
        else:
            phrase += ' ' + t1
    return phrase.capitalize()

Суть метода в том, что мы последовательно выбираем наиболее вероятные слова или знаки препинания до тех пор, пока не встречаем признак начала следующей фразы (символа $). Первое слово выбирается как наиболее вероятное для начала предложения из набора model['$', '$'].

Здесь необходимо отметить важный момент. Словарь model для каждой пары слов содержит список пар (слово, вероятность). Нам же необходимо выбрать из этого набора только одно слово. Вариант "в лоб" - выбрать слово с максимальной вероятностью. Но тогда все сгенерированные фразы были бы похожи друг на друга. Более подходящий способ способ - выбирать слова с некой хаотичностью, которая бы зависела от вероятности слова (мы же не хотим чтобы наши фразы состояли из редких сочетаний). Это и делает метод unirand, который возвращает случайное слово с вероятностью, равной вероятности данного слова в зависимости от двух предыдущих.

Итого, полный код нашего генератора выглядит следующим образом:

#!/usr/bin/env python
# -*- coding: utf-8 -*-

import re
from random import uniform
from collections import defaultdict

r_alphabet = re.compile(u'[а-яА-Я0-9-]+|[.,:;?!]+')

def gen_lines(corpus):
    data = open(corpus)
    for line in data:
        yield line.decode('utf-8').lower()

def gen_tokens(lines):
    for line in lines:
        for token in r_alphabet.findall(line):
            yield token

def gen_trigrams(tokens):
    t0, t1 = '$', '$'
    for t2 in tokens:
        yield t0, t1, t2
        if t2 in '.!?':
            yield t1, t2, '$'
            yield t2, '$','$'
            t0, t1 = '$', '$'
        else:
            t0, t1 = t1, t2

def train(corpus):
    lines = gen_lines(corpus)
    tokens = gen_tokens(lines)
    trigrams = gen_trigrams(tokens)

    bi, tri = defaultdict(lambda: 0.0), defaultdict(lambda: 0.0)

    for t0, t1, t2 in trigrams:
        bi[t0, t1] += 1
        tri[t0, t1, t2] += 1

    model = {}
    for (t0, t1, t2), freq in tri.iteritems():
        if (t0, t1) in model:
            model[t0, t1].append((t2, freq/bi[t0, t1]))
        else:
            model[t0, t1] = [(t2, freq/bi[t0, t1])]
    return model

def generate_sentence(model):
    phrase = ''
    t0, t1 = '$', '$'
    while 1:
        t0, t1 = t1, unirand(model[t0, t1])
        if t1 == '$': break
        if t1 in ('.!?,;:') or t0 == '$':
            phrase += t1
        else:
            phrase += ' ' + t1
    return phrase.capitalize()

def unirand(seq):
    sum_, freq_ = 0, 0
    for item, freq in seq:
        sum_ += freq
    rnd = uniform(0, sum_)
    for token, freq in seq:
        freq_ += freq
        if rnd < freq_:
            return token

if __name__ == '__main__':
    model = train('tolstoy.txt')
    for i in range(10):
        print generate_sentence(model),

Отлично, вы добрались до конца! Завидую вашему терпению :).

Почему триграммы

Триграммная модель выбрана для простоты и наглядности. Биграммы давали бы плохой результата, в то время как 4-граммы требовали бы существенно больше ресурсов. В любом случае, довольно просто расширить данный алгоритм для обработки общего случая N-грамм.