Char-RNN in PyTorch

char_rnn.py

import torch
import torch.optim as optim
import torch.nn as nn
import torch.nn.functional as F

class CharRNN(nn.Module):
    def __init__(self, vocab_size, model, hidden_size=256, num_layers=3):
        super().__init__()
        self.model = model
        self.num_layers = num_layers
        self.hidden_size = hidden_size
        self.embed = nn.Embedding(vocab_size, 128)
        if model == 'lstm':
            self.rnn = nn.LSTM(128, hidden_size, num_layers, dropout=0.2)
        elif model == 'gru':
            self.rnn = nn.GRU(128, hidden_size, num_layers, dropout=0.2)
        else:
            self.rnn = nn.RNN(128, hidden_size, num_layers, dropout=0.2)
        self.dense = nn.Linear(256, vocab_size)

    def forward(self, input, hidden):
        # input has shape (seq_len, batch)
        x = self.embed(input) # shape (seq_len, batch, 128)
        x, hidden = self.rnn(x, hidden) # shape (seq_len, batch, hidden_size)
        x = F.log_softmax(self.dense(x), dim=-1) # shape (seq_len, batch, vocab_size)
        return x, hidden

    def init_hidden(self, batch_size, **kw):
        if self.model != 'lstm':
            return torch.zeros((self.num_layers, batch_size, self.hidden_size), **kw)
        return (torch.zeros((self.num_layers, batch_size, self.hidden_size), **kw),
                torch.zeros((self.num_layers, batch_size, self.hidden_size), **kw))


class TextDataset:
    def __init__(self, filename):
        with open(filename, 'r', encoding='utf-8') as f:
            self.text = f.read()
        self.chars = list(set(self.text))
        self.char_to_idx = {c: i for i, c in enumerate(self.chars)}
        self.idx_to_char = {i: c for i, c in enumerate(self.chars)}
        self.vocab_size = len(self.chars)

    def get_segment(self, idx, length):
        return torch.tensor([self.char_to_idx[c] for c in self.text[idx : idx + length]])

    def batch_generator(self, batch_size, seq_length):
        length = len(self.text)
        batch_chars = length // batch_size
        for start in range(0, batch_chars - seq_length, seq_length):
            x = torch.zeros((seq_length, batch_size), dtype=torch.long)
            y = torch.zeros((seq_length, batch_size), dtype=torch.long)
            for batch_idx in range(0, batch_size):
                x[:, batch_idx] = self.get_segment(batch_chars * batch_idx + start, seq_length)
                y[:, batch_idx] = self.get_segment(batch_chars * batch_idx + start + 1, seq_length)
            yield x, y

Training.ipynb