# Create a new model to update the embeddings according to the requirement
class Modeler(nn.Module):
    
    def __init__(self, embed, vocab_size, embed_dim, keyword):
        super(Modeler, self).__init__()
        
        self.embeddings = nn.Embedding(vocab_size, embed_dim)
        self.embeddings.weight.data.copy_(torch.from_numpy(embed))
        self.embeddings.weight.requires_grad = True
        self.keyword = keyword
        self.linear1 = nn.Linear(embed_dim, 128)
        self.linear2 = nn.Linear(128, 1)
        
    def forward(self, input):
        embed = self.embeddings(input)
        k = embed.dot(self.keyword)
        out = F.relu(self.linear1(embed))
        out = self.linear2(out)
        return out
    
losses = []
loss_function = nn.MSELoss()
lookup_tensor = torch.LongTensor([word_to_ix['domain']])
embed = model.embeddings(Variable(lookup_tensor))
n_model = Modeler(model.embeddings.weight.data.numpy(), len(vocab),embedding_size, embed)
optimizer = optim.SGD(n_model.parameters(), lr=0.001)
batch_size = 1000
start = 0
end = batch_size 
for e in range(700):
     for epoch in range(int(len(data)/batch_size)):
        total_loss = torch.Tensor([0])
        for word, t in data[start:end]:
            word_id = word_to_ix[word]
            word_var = Variable(torch.LongTensor([word_id]))
            # Step 2. Recall that torch *accumulates* gradients. Before passing in a
            # new instance, you need to zero out the gradients from the old
            # instance
            n_model.zero_grad()            # Step 3. Run the forward pass, getting log probabilities over next
            # words
            res = n_model(word_var)
            # Step 4. Compute your loss function. (Again, Torch wants the target
            # word wrapped in a variable)
            t = Variable(torch.FloatTensor([t]))
            loss = loss_function(res, t)
            print(loss, t)
            # Step 5. Do the backward pass and update the gradient
            loss.backward()
            optimizer.step()
            total_loss += loss.data
            start += batch_size
            end += batch_size
        losses.append(total_loss)
print(np.mean(losses))