Neural Network Language Model Source Code

import torch
import torch.nn as nn
import torch.optim as optim

class NNLM(nn.Module):
    def __init__(self,input_layer, length, hidden_layer1, hidden_layer2, output_layer):
        super().__init__()
        self.hidden_layer_emb = hidden_layer1 * length
        self.C = nn.Embedding(input_layer, hidden_layer1)
        self.H = nn.Linear(self.hidden_layer_emb, hidden_layer2)
        self.d = nn.Parameter(torch.ones(hidden_layer2))
        self.U = nn.Linear(hidden_layer2, output_layer)
        self.W = nn.Linear(self.hidden_layer_emb, output_layer)
        self.b = nn.Parameter(torch.ones(output_layer))
        
        self.tanh = nn.Tanh()
        
    def forward(self, x):
        x = self.C(x)
        x = x.view(-1, self.hidden_layer_emb)
        res = self.H(x) + self.d
        res = self.U(self.tanh(res)) + self.W(x) + self.b
        return res

def tokenizer(sentence):
    dict_number = {}
    dict_word = {}
    for s in sentence:
        for word in s.split():
            if word not in dict_number.keys():
                dict_number[word] = len(dict_number)
                dict_word[len(dict_number)-1] = word
                
    return dict_number, dict_word

def make_data(sentence,dict_number):
    X = []
    Y = []
    for sen in sentence:
        s = sen.split()
        X.append([dict_number[w] for w in s[:-1]])
        Y.append(dict_number[s[-1]])
    return X,Y

sentences = [
    'i love you',
    'i hate them',
    'i hit ball',
    'i kill mouse',
    'you love me',
    'you hate them',
    'you kill mouse',
    'they like us',
    'they are bad'
]

dict_number, dict_word = tokenizer(sentences)
X,Y = make_data(sentences, dict_number)

model = NNLM(len(dict_number),2,3,3,len(dict_number))

loss_fn = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

epochs = 10000

X,Y = torch.LongTensor(X),torch.LongTensor(Y)


for epoch in range(epochs):
    Y_pred = model(X)
    loss = loss_fn(Y_pred,Y)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    if (epoch+1)%500==0:
        print(epoch+1, loss.float())

def test_model(model,dict_number,sentence):
    x = torch.LongTensor([[dict_number[i] for i in sentence.split()]])
    predict = model(x).data.max(1, keepdim=True)[1]
    print(sentence,dict_word[predict.squeeze().item()])

test_model(model,dict_number,'you love')
# >>> you love me