本人小菜鸟，也就是有点闲做了个写小说人工智能而已，RNN架构，欢迎大家自己跑一下（显存2G或没有独显的同志请做好迎接电脑爆炸的决心）恳请站长指导！

                                                  声明
本项目基于pytorch实现。可以自行更换训练文本，但请记住：训练文本一定一定得是TXT格式！！！
另外，本项目仅用作学习交流，严禁用于商用、参加任何形式以及团体组织的比赛，请使用者严格遵守Apache License Version 开源协议
作者不承担此程序造成的一切后果与责任。
版权归作者重庆市小学生胡哲涵所有，Github likehuiyuanai。
PS：最近喜欢看某科学的超电磁炮，所以实例代码用的是日本轻小说《魔法禁书目录》作为训练数据，暂拒绝公开此数据集（你可以自己整理，字数多亿些训练效果好点
本博客（帖子）的版权也同样归胡哲涵所有，严禁转发或参加比赛！

import torch
import torch.nn as nn
import numpy as np
import matplotlib.pyplot as plt
import time
from scipy.sparse import csr_matrix
from tensorboardX import SummaryWriter
%matplotlib inline

with open('./mfjsml.txt', 'r', encoding='utf-8') as f:
    data = f.readlines()

data=''.join(data)

print(data[:100])

chars = list(set(data))
data_size, vocab_size = len(data), len(chars)
print(f'data has {data_size} characters, {vocab_size} unique.')
char_to_ix = { ch:i for i,ch in enumerate(chars) }
ix_to_char = { i:ch for i,ch in enumerate(chars) }

X_train = csr_matrix((len(data), len(chars)), dtype=np.int)
char_id = np.array([chars.index(c) for c in data])
X_train[np.arange(len(data)), char_id] = 1

y_train = np.roll(char_id,-1)

X_train.shape

y_train.shape

def get_batch(X_train, y_train, seq_length):
    '''Return a training batch with certain number of X and y pairs.'''
    X = X_train
    #X = torch.from_numpy(X_train).float()
    y = torch.from_numpy(y_train).long()
    for i in range(0, len(y), seq_length):   
        id_stop = i+seq_length if i+seq_length < len(y) else len(y)
        yield([torch.from_numpy(X[i:id_stop].toarray().astype(np.float32)), 
               y[i:id_stop]])

def sample_chars(rnn, X_seed, h_prev, length=20):
    '''Generate text using trained model'''
    X_next = X_seed
    results = []
    with torch.no_grad():
        for i in range(length):        
            y_score, h_prev = rnn(X_next.view(1,1,-1), h_prev)
            y_prob = nn.Softmax(0)(y_score.view(-1)).detach().numpy()
            y_pred = np.random.choice(chars,1, p=y_prob).item()
            results.append(y_pred)
            X_next = torch.zeros_like(X_seed)
            X_next[chars.index(y_pred)] = 1
    return ''.join(results)

class nn_LSTM(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super().__init__()
        self.hidden_size = hidden_size
        self.lstm = nn.LSTM(input_size, hidden_size)
        self.out = nn.Linear(hidden_size, output_size)
        
    def forward(self, X, hidden):
        _, hidden = self.lstm(X, hidden)
        output = self.out(hidden[0])
        return output, hidden
    
    def initHidden(self):
        return (torch.zeros(1, 1, self.hidden_size),
                torch.zeros(1, 1, self.hidden_size)
               )

hidden_size = 256
seq_length = 25

rnn = nn_LSTM(vocab_size, hidden_size, vocab_size)

loss_fn = nn.CrossEntropyLoss()

optimizer = torch.optim.Adam(rnn.parameters(), lr=0.005)

def train(X_batch, y_batch):
    h_prev = rnn.initHidden()
    optimizer.zero_grad()
    batch_loss = torch.tensor(0, dtype=torch.float)
    
    for i in range(len(X_batch)):
        y_score, h_prev = rnn(X_batch[i].view(1,1,-1), h_prev)
        loss = loss_fn(y_score.view(1,-1), y_batch[i].view(1))
        batch_loss += loss
    batch_loss.backward()
    optimizer.step()

    return y_score, batch_loss/len(X_batch)

writer = SummaryWriter(f'logs/lstm1_{time.strftime("%Y%m%d-%H%M%S")}')

all_losses = []
print_every = 100
for epoch in range(20):    
    for batch in get_batch(X_train, y_train, seq_length):
        X_batch, y_batch = batch
        _, batch_loss = train(X_batch, y_batch)
        all_losses.append(batch_loss.item())
        if len(all_losses)%print_every==1:
            print(f'----\nRunning Avg Loss:{np.mean(all_losses[-print_every:])} at iter: {len(all_losses)}\n----')
            # log to tensorboard every X iterations. Can be removed if Tensorboard is not installed.
            writer.add_scalar('loss', np.mean(all_losses[-100:]), len(all_losses))
            # generate text every X iterations
            print(sample_chars(rnn, X_batch[0], rnn.initHidden(), 200))

print(sample_chars(rnn, X_batch[20], rnn.initHidden(), 200))

torch.save(rnn.state_dict(), 'shediao.pth')