Memory Networks

from babi_loader import BabiDataset, pad_collate

from torch.utils.data import DataLoader
# There are 20 tasks, we should control which task we would like to load
task_id = 1
dataset = BabiDataset( task_id )
dataloader = DataLoader(dataset, batch_size=4, shuffle=False, collate_fn=pad_collate)

contexts, questions, answers = next(iter(dataloader))
print(contexts.shape)
print(questions.shape)
print(answers.shape)
print(contexts)
print(questions)
print(answers)

references = dataset.QA.IVOCAB
print(references)

def interpret_indexed_tensor(contexts, questions, answers):
    for n, data in enumerate(zip(contexts, questions, answers)):
        context  = data[0]
        question = data[1]
        answer   = data[2]
        
        print(n)
        for i, sentence in enumerate(context):
            s = ' '.join([references[elem.item()] for elem in sentence])
            print(f'{i}th sentence: {s}')
        s = ' '.join([references[elem.item()] for elem in question])
        print(f'question: {s}')
        s = references[answer.item()]
        print(f'answer:   {s}')

interpret_indexed_tensor(contexts, questions, answers)

import os
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.nn.init as init
from torch.autograd import Variable

class InputModule(nn.Module):
    def __init__(self, vocab_size, hidden_size):
        super(InputModule, self).__init__()
        self.hidden_size = hidden_size
        
        ##################################################################################
        #
        # nn.GRU use the bidirectional parameter to control the type of GRU
        #
        ##################################################################################
        self.gru = nn.GRU(hidden_size, hidden_size, bidirectional=False, batch_first=True)
        self.dropout = nn.Dropout(0.1)

    # we should all the word_embedding the same, while we using the different module
    def forward(self, contexts, word_embedding):
        '''
        contexts.size() -> (#batch, #sentence, #token)
        word_embedding() -> (#batch, #sentence x #token, #embedding)
        position_encoding() -> (#batch, #sentence, #embedding)
        facts.size() -> (#batch, #sentence, #hidden = #embedding)
        '''
        batch_num, sen_num, token_num = contexts.size()

        contexts = contexts.view(batch_num, -1)
        contexts = word_embedding(contexts)

        contexts = contexts.view(batch_num, sen_num, token_num, -1)
        contexts = self.position_encoding(contexts)
        contexts = self.dropout(contexts)

        #########################################################################
        #
        # if you change the gru type, you should also change the initial hidden state
        # as bidirectional gru's shape is (2, *, *), while normal gru just need( 1, *, *)
        #
        #########################################################################
        h0 = Variable(torch.zeros(1, batch_num, self.hidden_size).cuda())
        facts, hdn = self.gru(contexts, h0)
        #########################################################################
        #
        # if you use bi-directional GRU, you should fusion the output,
        # if you use normal GRU, commont the following code. 
        # acconding to the equation (6)
        #
        #########################################################################
#         facts = facts[:, :, :self.hidden_size] + facts[:, :, self.hidden_size:]
        return facts
    
    def position_encoding(self, embedded_sentence):
        '''
        embedded_sentence.size() -> (#batch, #sentence, #token, #embedding)
        l.size() -> (#sentence, #embedding)
        output.size() -> (#batch, #sentence, #embedding)
        '''
        _, _, slen, elen = embedded_sentence.size()

        l = [[(1 - s/(slen-1)) - (e/(elen-1)) * (1 - 2*s/(slen-1)) for e in range(elen)] for s in range(slen)]
        l = torch.FloatTensor(l)
        l = l.unsqueeze(0) # for #batch
        l = l.unsqueeze(1) # for #sen
        l = l.expand_as(embedded_sentence)
        weighted = embedded_sentence * Variable(l.cuda())
        return torch.sum(weighted, dim=2).squeeze(2) # sum with tokens

class QuestionModule(nn.Module):
    def __init__(self, vocab_size, hidden_size):
        super(QuestionModule, self).__init__()
        self.gru = nn.GRU(hidden_size, hidden_size, batch_first=True)

    def forward(self, questions, word_embedding):
        '''
        questions.size() -> (#batch, #token)
        word_embedding() -> (#batch, #token, #embedding)
        gru() -> (1, #batch, #hidden)
        '''
        questions = word_embedding(questions)
        _, questions = self.gru(questions)
        questions = questions.transpose(0, 1)
        return questions

class AttentionGRUCell(nn.Module):
    def __init__(self, input_size, hidden_size):
        super(AttentionGRUCell, self).__init__()
        self.hidden_size = hidden_size
        
        self.Wr = nn.Linear(input_size, hidden_size)
        self.Ur = nn.Linear(hidden_size, hidden_size)
        self.W = nn.Linear(input_size, hidden_size)
        self.U = nn.Linear(hidden_size, hidden_size)

    def forward(self, fact, C, g):
        r = torch.sigmoid(self.Wr(fact) + self.Ur(C))
        h_tilda = torch.tanh(self.W(fact) + r * self.U(C))
        g = g.unsqueeze(1).expand_as(h_tilda)
        h = g * h_tilda + (1 - g) * C
        return h

class AttentionGRU(nn.Module):
    def __init__(self, input_size, hidden_size):
        super(AttentionGRU, self).__init__()
        self.hidden_size = hidden_size
        self.AGRUCell = AttentionGRUCell(input_size, hidden_size)

    def forward(self, facts, G):
        batch_num, sen_num, embedding_size = facts.size()
        C = Variable(torch.zeros(self.hidden_size)).cuda()
        for sid in range(sen_num):
            fact = facts[:, sid, :]
            g = G[:, sid]
            if sid == 0:
                C = C.unsqueeze(0).expand_as(fact)
            C = self.AGRUCell(fact, C, g)
        return C

class EpisodicMemory(nn.Module):
    def __init__(self, hidden_size):
        super(EpisodicMemory, self).__init__()
        self.AGRU = AttentionGRU(hidden_size, hidden_size)
        self.z1 = nn.Linear(4 * hidden_size, hidden_size)
        self.z2 = nn.Linear(hidden_size, 1)
        self.next_mem = nn.Linear(3 * hidden_size, hidden_size)

    def make_interaction(self, facts, questions, prevM):
        '''
        facts.size() -> (#batch, #sentence, #hidden = #embedding)
        questions.size() -> (#batch, 1, #hidden)
        prevM.size() -> (#batch, #sentence = 1, #hidden = #embedding)
        z.size() -> (#batch, #sentence, 4 x #embedding)
        G.size() -> (#batch, #sentence)
        '''
        batch_num, sen_num, embedding_size = facts.size()
        questions = questions.expand_as(facts)
        prevM = prevM.expand_as(facts)

        z = torch.cat([
            facts * questions,
            facts * prevM,
            torch.abs(facts - questions),
            torch.abs(facts - prevM)
        ], dim=2)

        z = z.view(-1, 4 * embedding_size)

        G = torch.tanh(self.z1(z))
        G = self.z2(G)
        G = G.view(batch_num, -1)
        G = F.softmax(G, dim=1)

        return G
    
    def forward(self, facts, questions, prevM):
        '''
        facts.size() -> (#batch, #sentence, #hidden = #embedding)
        questions.size() -> (#batch, #sentence = 1, #hidden)
        prevM.size() -> (#batch, #sentence = 1, #hidden = #embedding)
        G.size() -> (#batch, #sentence)
        C.size() -> (#batch, #hidden)
        concat.size() -> (#batch, 3 x #embedding)
        '''
        G = self.make_interaction(facts, questions, prevM)
        C = self.AGRU(facts, G)
        concat = torch.cat([prevM.squeeze(1), C, questions.squeeze(1)], dim=1)
        next_mem = F.relu(self.next_mem(concat))
        next_mem = next_mem.unsqueeze(1)
        return next_mem

class AnswerModule(nn.Module):
    def __init__(self, vocab_size, hidden_size):
        super(AnswerModule, self).__init__()
        self.z = nn.Linear(2 * hidden_size, vocab_size)
        self.dropout = nn.Dropout(0.1)

    def forward(self, M, questions):
        M = self.dropout(M)
        concat = torch.cat([M, questions], dim=2).squeeze(1)
        z = self.z(concat)
        return z

class DMNPlus(nn.Module):
    '''
    This class combine all the module above. The data flow is showed as the above image.
    '''
    def __init__(self, hidden_size, vocab_size, num_hop=3, qa=None):
        super(DMNPlus, self).__init__()
        self.num_hop = num_hop
        self.qa = qa
        self.word_embedding = nn.Embedding(vocab_size, hidden_size, padding_idx=0, sparse=True).cuda()
        self.criterion = nn.CrossEntropyLoss(reduction='sum')

        self.input_module = InputModule(vocab_size, hidden_size)
        self.question_module = QuestionModule(vocab_size, hidden_size)
        self.memory = EpisodicMemory(hidden_size)
        self.answer_module = AnswerModule(vocab_size, hidden_size)

    def forward(self, contexts, questions):
        '''
        contexts.size() -> (#batch, #sentence, #token) -> (#batch, #sentence, #hidden = #embedding)
        questions.size() -> (#batch, #token) -> (#batch, 1, #hidden)
        '''
        facts = self.input_module(contexts, self.word_embedding)
        questions = self.question_module(questions, self.word_embedding)
        M = questions
        for hop in range(self.num_hop):
            M = self.memory(facts, questions, M)
        preds = self.answer_module(M, questions)
        return preds

    # train the index into a word, it's similar to the part 1 data explore
    def interpret_indexed_tensor(self, var):
        if len(var.size()) == 3:
            # var -> n x #sen x #token
            for n, sentences in enumerate(var):
                for i, sentence in enumerate(sentences):
                    s = ' '.join([self.qa.IVOCAB[elem.data[0]] for elem in sentence])
                    print(f'{n}th of batch, {i}th sentence, {s}')
        elif len(var.size()) == 2:
            # var -> n x #token
            for n, sentence in enumerate(var):
                s = ' '.join([self.qa.IVOCAB[elem.data[0]] for elem in sentence])
                print(f'{n}th of batch, {s}')
        elif len(var.size()) == 1:
            # var -> n (one token per batch)
            for n, token in enumerate(var):
                s = self.qa.IVOCAB[token.data[0]]
                print(f'{n}th of batch, {s}')

    # calculate the loss of the network
    def get_loss(self, contexts, questions, targets):
        output = self.forward(contexts, questions)
        loss = self.criterion(output, targets)
        reg_loss = 0
        for param in self.parameters():
            reg_loss += 0.001 * torch.sum(param * param)
        preds = F.softmax(output, dim=1)
        _, pred_ids = torch.max(preds, dim=1)
        corrects = (pred_ids.data == answers.cuda().data)
        acc = torch.mean(corrects.float())
        return loss + reg_loss, acc

task_id = 5

epochs = 3

if task_id in range(1, 21):
#def train(task_id):
    dset = BabiDataset(task_id)
    vocab_size = len(dset.QA.VOCAB)
    hidden_size = 80
    

    model = DMNPlus(hidden_size, vocab_size, num_hop=3, qa=dset.QA)
    model.cuda()
    
    best_acc = 0
    optim = torch.optim.Adam(model.parameters())

    for epoch in range(epochs):
        dset.set_mode('train')
        train_loader = DataLoader(
                dset, batch_size=32, shuffle=True, collate_fn=pad_collate
                )

        model.train()
        total_acc = 0
        cnt = 0
        for batch_idx, data in enumerate(train_loader):
            contexts, questions, answers = data
            batch_size = contexts.size()[0]
            contexts  = Variable(contexts.long().cuda())
            questions = Variable(questions.long().cuda())
            answers   = Variable(answers.cuda())

            loss, acc = model.get_loss(contexts, questions, answers)
            optim.zero_grad()
            loss.backward()
            optim.step()
            
            total_acc += acc * batch_size
            cnt += batch_size
            if batch_idx % 20 == 0:
                print(f'[Task {task_id}, Epoch {epoch}] [Training] loss : {loss.item(): {10}.{8}}, acc : {total_acc / cnt: {5}.{4}}, batch_idx : {batch_idx}')

dset.set_mode('test')
test_loader = DataLoader(
    dset, batch_size=100, shuffle=False, collate_fn=pad_collate
    )
test_acc = 0
cnt = 0

model.eval()
for batch_idx, data in enumerate(test_loader):
    contexts, questions, answers = data
    batch_size = contexts.size()[0]
    contexts = Variable(contexts.long().cuda())
    questions = Variable(questions.long().cuda())
    answers = Variable(answers.cuda())
    
    _, acc = model.get_loss(contexts, questions, answers)
    test_acc += acc * batch_size
    cnt += batch_size
print(f'Task {task_id}, Epoch {epoch}] [Test] Accuracy : {test_acc / cnt: {5}.{4}}')

def interpret_indexed_tensor(references, contexts, questions, answers, predictions):
    for n, data in enumerate(zip(contexts, questions, answers, predictions)):
        context  = data[0]
        question = data[1]
        answer   = data[2]
        predict  = data[3]
        
        print(n)
        for i, sentence in enumerate(context):
            s = ' '.join([references[elem.item()] for elem in sentence])
            print(f'{i}th sentence: {s}')
        q = ' '.join([references[elem.item()] for elem in question])
        print(f'question: {q}')
        a = references[answer.item()]
        print(f'answer:   {a}')
        p = references[predict.argmax().item()]
        print(f'predict:  {p}')

dset.set_mode('test')
dataloader = DataLoader(dset, batch_size=4, shuffle=False, collate_fn=pad_collate)
contexts, questions, answers = next(iter(dataloader))

contexts  = Variable(contexts.long().cuda())
questions = Variable(questions.long().cuda())

# prediction
model.eval()
predicts = model(contexts, questions)

references = dset.QA.IVOCAB
#contexts, questions, answers = next(iter(dataloader))
#interpret_indexed_tensor(contexts, questions, answers)
interpret_indexed_tensor(references, contexts, questions, answers, predicts)

dataset = BabiDataset( task_id )

# dataset.set_mode("train")
dataset.set_mode("test")

print(len(dataset))