diff --git a/rnn-poetry/model.py b/rnn-poetry/model.py new file mode 100644 index 0000000..05e8d16 --- /dev/null +++ b/rnn-poetry/model.py @@ -0,0 +1,58 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F + +batch_size=128 +embed_size=128 +hidden_dims=256 + +def generate_poetry(model,word2ix,ix2word,device,begin,sent_len=4): + start_idx=[word2ix['[']] + end_word='' + lens=0 + hidden = None + ret='' + data_ = torch.tensor([start_idx], device=device).long() + output, hidden = model(data_, hidden) + start_idx=[word2ix[begin]] + ret+=begin + while end_word!=']' and len(ret)<100: + data_ = torch.tensor([start_idx],device=device).long() + # print("data size",data_.size()) + output, hidden = model(data_, hidden) + # print("output size", output.size()) + ouput_idx=output.view(-1).argmax().cpu() + # print('ouput_idx',ouput_idx) + # print('ouput_idx', ouput_idx.item()) + ouput_idx=ouput_idx.item() + start_idx=[ouput_idx] + end_word=ix2word[ouput_idx] + ret+=end_word + return ret + +class RNNModel(nn.Module): + def __init__(self, vocab_size, embedding_dim, hidden_dim): + super(RNNModel, self).__init__() + self.hidden_dim = hidden_dim + self.embeddings = nn.Embedding(vocab_size, embedding_dim) + self.lstm = nn.LSTM(embedding_dim, self.hidden_dim, num_layers=2) + self.linear1 = nn.Linear(self.hidden_dim, vocab_size) + + + + def forward(self, x, hidden=None): + seq_len, batch_size = x.size() + + + # size: (seq_len,batch_size,embeding_dim) + embeds = self.embeddings(x) + # output size: (seq_len,batch_size,hidden_dim) + if hidden is None: + output, hidden = self.lstm(embeds) + else: + h_0, c_0 = hidden + output, hidden = self.lstm(embeds, (h_0, c_0)) + + # size: (seq_len*batch_size,vocab_size) + output = self.linear1(output.view(seq_len * batch_size, -1)) + return output, hidden