Как научить нейросеть генерировать текст с помощью LSTM в PyTorch

Статья написана по мотивам очередной домашки в MADE, где мы учили нейронку писать стихи с помощью рекурентных сетей.

Для начала, договоримся, что будем делать «глупую» нейросеть, которая не разбирается в языке, построении фраз, предложений, смыслах, а просто учится предсказывать следующий символ по предыдущему тексту.

Из плюсов такого подхода — не нужно париться о пунктуации, больших и маленьких буквах, нейронка их расставит сама. Из минусов — такая нейросеть «туповата».

Так же можно генерировать любые последовательности. К примеру, музыку. MIDI конвертируем в формат ABC, получаем текст (точнее, последовательность символов). На этом тексте учимся, генерируем свой, превращаем его обратно в MIDI и получаем «нейромузыку».

Если генерировать тексты, будет получаться что-то вроде:

Потом ахали на наших специалистов: они, не имея вакцины, не зная, как мы действуем сейчас: малейшее подозрение — мы проверяем человека, надо на карантин — пожалуйста, давайте на карантин. Именно точечно, профессионально надо действовать, а не территориально не сможем, это невозможно. Особенно там, но полоза. Потому что они живут, как обычно (находятся не только в каких-то все обсуждать, но и проблема. Надо «желтые жилеты» вернуть с улиц. У кого-то выборы. У кого-то еще чего-то. И каждый начинает вертерах на неменех, не дали было спросы.

Нейросеть, обученная на изречениях бессменного лидера РБ

А иногда выходит что-то совсем прекрасное:

Стоит ли любить?
Стоит ли любить?
Стоит ли любить?
Стоит и приплыли.

Нейросеть, обученная на хокку

Практическое применение этому мне сложно придумать. Скорее, это — про побаловаться. Ну и сотворить что-то забавное. Fake news этом не сделаешь. Точнее, сделаешь, но они будут наредкость забавны.

И важный момент. Если что-то непонятно или просто неинтересно, а попробовать генерировать текст хочется, то достаточно просто перейти по ссылке в Google Colab в конце статьи, закинуть туда свой файлик с текстом и запустить все ячейки с кодом.

Итак, поехали.

1. Импортируем нужные библиотеки

Из внешних библиотек нам понадобятся только numpy и pytorch:

from collections import Counter

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np

2. Готовим данные для сети

Токенизируем текст. Превращаем его в индексы. Т.е. достаем все уникальные символы: буквы, пробелы, знаки препинания в тексте и каждому символу присваиваем число. Например:

• ‘a’ => 1
• ‘Щ’ => 2
• ‘!’ => 3
• …

Это будем называть словарем. А числа — индексами.

Делаем прямой и обратный (индексы в символы) словарь. А потом проходимся по всему тексту и превращаем с помощью составленного нами словаря каждый символ в индекс.

TRAIN_TEXT_FILE_PATH = 'train_text.txt'

with open(TRAIN_TEXT_FILE_PATH) as text_file:
    text_sample = text_file.readlines()
text_sample = ' '.join(text_sample)

def text_to_seq(text_sample):
    char_counts = Counter(text_sample)
    char_counts = sorted(char_counts.items(), key = lambda x: x[1], reverse=True)

    sorted_chars = [char for char, _ in char_counts]
    print(sorted_chars)
    char_to_idx = {char: index for index, char in enumerate(sorted_chars)}
    idx_to_char = {v: k for k, v in char_to_idx.items()}
    sequence = np.array([char_to_idx[char] for char in text_sample])
    
    return sequence, char_to_idx, idx_to_char

sequence, char_to_idx, idx_to_char = text_to_seq(text_sample)

3. Генерируем батчи из текста

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

Будем генерировать сразу обучающую выборку (то, на чем будем учить сеть) и таргет для нее. Таргет (правильные ответы для нейросети) — это просто сдвинутый на один символ вперед текст.

Условно:

• Обучающая выборка: Привет как дел
• Таргет: ривет как дела

Размерность тензора батча: [BATCH_SIZE x SEQ_LEN x 1]

SEQ_LEN = 256
BATCH_SIZE = 16

def get_batch(sequence):
    trains = []
    targets = []
    for _ in range(BATCH_SIZE):
        batch_start = np.random.randint(0, len(sequence) - SEQ_LEN)
        chunk = sequence[batch_start: batch_start + SEQ_LEN]
        train = torch.LongTensor(chunk[:-1]).view(-1, 1)
        target = torch.LongTensor(chunk[1:]).view(-1, 1)
        trains.append(train)
        targets.append(target)
    return torch.stack(trains, dim=0), torch.stack(targets, dim=0)

4. Пишем функцию, которая генерирует текст

Теперь напишем функцию, которая предсказывает текст с помощью нашей обученной нейросети. Это удобно сделать заранее, чтобы смотреть, что генерирует сеть во время обучения.

Сеть предсказывает нам вероятности следующей буквы, и мы с помощью этих вероятностей достаем случайно по одной букве. Если повторить операцию 1000 раз, получим текст из 1000 символов.

Параметр start_text нам нужен, чтобы было что-то, для чего предсказывать следующий символ. У нас этот символ по умолчанию — пробел, и задача сети сначала — предсказать следующий символ после пробела. Потом — следующий после этих 2-х символов. И т.д.

Параметр temp — это уровень «случайности» генерируемого текста. Так называемая «температура» с отсылкой к понятию «энтропии». То, на что делим логиты в softmax. Поставим высокую — вероятность каждой буквы будет почти одинакова и текст превратится в случайную белиберду. Поставим низкую — каждый раз будем предсказывать одно и то же и можем зациклиться на одной фразе.

def evaluate(model, char_to_idx, idx_to_char, start_text=' ', prediction_len=200, temp=0.3):
    hidden = model.init_hidden()
    idx_input = [char_to_idx[char] for char in start_text]
    train = torch.LongTensor(idx_input).view(-1, 1, 1).to(device)
    predicted_text = start_text
    
    _, hidden = model(train, hidden)
        
    inp = train[-1].view(-1, 1, 1)
    
    for i in range(prediction_len):
        output, hidden = model(inp.to(device), hidden)
        output_logits = output.cpu().data.view(-1)
        p_next = F.softmax(output_logits / temp, dim=-1).detach().cpu().data.numpy()        
        top_index = np.random.choice(len(char_to_idx), p=p_next)
        inp = torch.LongTensor([top_index]).view(-1, 1, 1).to(device)
        predicted_char = idx_to_char[top_index]
        predicted_text += predicted_char
    
    return predicted_text

5. Создаем класс нашей нейросети

И наконец наша маленькая и уютная нейросеть. Она работает так:

1. Превращаем каждый символ на входе сети в вектор (так называемный эмбеддинг).
2. Скармливаем эти векторы нашему LSTM слою. У этого слоя есть особенность: он работает не независимо для каждого символа, а помнит, что к нему раньше приходило на вход. Притом, помнит не все: ненужное он умеет забывать. Такие слои называют рекуррентными и часто используют при работе с последовательностями.
3. Выходы из LSTM слоя пропускаем через Dropout. Этот слой «мешает» сети учиться, чтобы ей сложнее было выучить весть текст.
4. Дальше отправляем выход из Dropout на линейный слой размерности словаря, чтобы на выходе получить столько чисел, сколько у нас символов в словаре. Потом мы этот вектор чисел будем превращать в «вероятности» каждого символа с помощью функции softmax.

class TextRNN(nn.Module):
    
    def __init__(self, input_size, hidden_size, embedding_size, n_layers=1):
        super(TextRNN, self).__init__()
        
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.embedding_size = embedding_size
        self.n_layers = n_layers

        self.encoder = nn.Embedding(self.input_size, self.embedding_size)
        self.lstm = nn.LSTM(self.embedding_size, self.hidden_size, self.n_layers)
        self.dropout = nn.Dropout(0.2)
        self.fc = nn.Linear(self.hidden_size, self.input_size)
        
    def forward(self, x, hidden):
        x = self.encoder(x).squeeze(2)
        out, (ht1, ct1) = self.lstm(x, hidden)
        out = self.dropout(out)
        x = self.fc(out)
        return x, (ht1, ct1)
    
    def init_hidden(self, batch_size=1):
        return (torch.zeros(self.n_layers, batch_size, self.hidden_size, requires_grad=True).to(device),
               torch.zeros(self.n_layers, batch_size, self.hidden_size, requires_grad=True).to(device))

6. Создаем нейросеть и обучаем ее

Теперь создаем нейросеть и обучаем ее. LSTM блок принимает немного другой формат батча:

[SEQ_LEN x BATCH_SIZE x 1], поэтому делаем permute для тензоров train и target, чтобы поменять 0 и 1 размерность местами.

Параметры нейросети, которые может понадобиться подкрутить:

• hidden_size — влияет на сложность сети. Стоит повышать для текстов большого размера. Если выставить большое значение для текста маленького размера, то сеть просто выучит весь текст и будет генерировать его же.
• n_layers — опять же, влияет на сложность сети. Грубо говоря, позволяет делать несколько LSTM слоев подряд просто меняя эту цифру.
• embedding_size — размер обучаемого эмбеддинга. Можно выставить в несколько раз меньше размера словаря (числа уникальных символов в тексте) или примерно такой же. Больше — нет смысла.

Дальше — стандартный для PyTorch цикл обучения нейросети: выбираем функцию потерь, оптимизатор и настраиваем расписание, по которому меняем шаг оптимизатора. В нашем случае снижаем шаг в 2 раза, если ошибка (loss) не падает 5 шагов подряд.

Если очень грубо, то здесь мы много раз подаем в нейросеть разные кусочки текста и учим ее делать все меньше и меньше ошибок, когда она предсказывает следующую букву по предыдущему тексту.

device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
model = TextRNN(input_size=len(idx_to_char), hidden_size=128, embedding_size=128, n_layers=2)
model.to(device)

criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-2, amsgrad=True)
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
    optimizer, 
    patience=5, 
    verbose=True, 
    factor=0.5
)

n_epochs = 50000
loss_avg = []

for epoch in range(n_epochs):
    model.train()
    train, target = get_batch(sequence)
    train = train.permute(1, 0, 2).to(device)
    target = target.permute(1, 0, 2).to(device)
    hidden = model.init_hidden(BATCH_SIZE)

    output, hidden = model(train, hidden)
    loss = criterion(output.permute(1, 2, 0), target.squeeze(-1).permute(1, 0))
    
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()
    
    loss_avg.append(loss.item())
    if len(loss_avg) >= 50:
        mean_loss = np.mean(loss_avg)
        print(f'Loss: {mean_loss}')
        scheduler.step(mean_loss)
        loss_avg = []
        model.eval()
        predicted_text = evaluate(model, char_to_idx, idx_to_char)
        print(predicted_text)

При обучении должно получиться на выходе что-то вроде:

Loss: 3.0337722063064576
В то не не посто на на не не на на не на помотророне на онот потарити не сем посто нат на рото сто на но на то тосте на посто нал на на на на росмоно посто но это воме на не вотато не не но на не нам на
Loss: 2.4729482173919677
В поторовововать оточели не не тора на на на на подовать не то не на подени не на вобомать не не прорус не сто на на сторы домем водем не сто не это соторовоста вомать на на проденить не бостова сто пор
Loss: 2.2126503944396974
В это надо поторовать на содолько не полодали не то начения. Нако не надо на работать на на полили на на наши наши подули сотовать на нало обранить на на постовать не наши на поставил в ирания стравом м
Loss: 2.028042938709259
В кому просторы на наши на комперами на только полоблимать должны поровать не том приратить подурать рас совеловали полжно от востание в мостам в сего принимали сотрать в нам нам от проблимать больно на
.......
.......
.......
Loss: 0.5439809119701385
В ибыратий от безработицы и доклады. Авторов мы сегодня по тому сли нас дороге, чтобы они смогли выйти и пережил уже многие психозом особенно надо выживать на эту тему в этот пережил уже многие психозы,
Loss: 0.5482879960536957
В ибывать на сегодняшний день отношение домерми вде выписов из меня частся с вами и серация — ни одного сказать, что мы сейчас делаем, они понятный их сполоси, которые основной пофтомие потери людям за 
Loss: 0.5576688635349274
В конятно к этой пришетов не проблема станок их работать, потому что это повторилось.
 Результатом стал устойчивый рост благосостояния населения. Мы точно дошло предложено. Политически запрещаю. Давайте

Можете остановить обучение в любой момент. К примеру, тогда, когда ошибка (Loss) перестанет снижаться.

7. Генерируем текст

Итак, наша сеть обучилась. Давайте что-нибудь сгенерируем:

model.eval()
print(evaluate(
    model, 
    char_to_idx, 
    idx_to_char, 
    temp=0.3, 
    prediction_len=1000, 
    start_text='. '
    )
)

Рекомендую поиграть с параметром temp и, конечно же, start_text. С помощью start_text можно попробовать «задать тему/направление» для генерируемого текста. Желательно, ту, которая есть в тексте, на котором сеть училась, конечно.

Надеюсь, все получилось и нигде ничего не свалилось.

Не хочу ничего делать, хочу сразу генерировать текст

Так тоже можно. Просто откройте jupyter notebook в Google Colab, скопируйте туда файлик с текстом и выполните все ячейки.

Открыть готовый код в Google Colab

После того, как открыли, нужно скопировать туда и файл с текстом, на котором будет тренироваться нейросеть.

Вот и все. Можно развлекаться. Просто запустите подряд все ячейки с кодом. Когда будете учить нейросеть (предпоследняя ячейка), можете не ждать, пока обучение дойдет до конца и остановить в любой момент.

И еще крайне рекомендую в настройках Colab «Runtime» => «Change runtime type» => «Hardware accelerator» выбрать «GPU». На видеокарте сеть учится в разы быстрее.

На всякий случай, выложил код на github.

Если есть какие-то вопросы, что-то не работает или вы нашли баг, пишите, постараюсь помочь.

P.S.: если получится сгенерировать что-то забавное, сбросьте это сюда в комментарии 🙂