!git clone https://github.com/haochen23/pytorch-nlp-classifier.git
%cd pytorch-nlp-classifier/

# Adopted from utils.py
def load_data(file_path):
    '''
    load and prepare dataset to training and validation iterator
    '''
    TEXT = data.Field(tokenize="spacy", batch_first=True, include_lengths=True)
    LABEL = data.LabelField(dtype=torch.float, batch_first=True)

    fields = [(None,None), ('text',TEXT), ('label',LABEL)]

    total_data = data.TabularDataset(path=file_path,
                                    format="csv", 
                                    fields=fields,
                                    skip_header=True)
    # split data
    train_data, valid_data = total_data.split(split_ratio=0.7, random_state = random.seed(seed))
    # initialize glove embeddings
    TEXT.build_vocab(train_data, min_freq=3, vectors="glove.6B.100d")
    LABEL.build_vocab(train_data)
    batch_size = config.BATCH_SIZE
    device = config.device
    train_iterator, valid_iterator = data.BucketIterator.splits(
        (train_data, valid_data),
        batch_size = batch_size,
        sort_key = lambda x: len(x.text),
        sort_within_batch = True,
        device=device
    )

    return train_iterator, valid_iterator, TEXT, LABEL

# Adopted from classifer.py
class classifer(nn.Module):
    #define all the layers used in the model
    def __init__(self, vocab_size, embedding_dim, hidden_dim, output_dim, n_layers, 
                 bidirectional, dropout):
        super().__init__()

        # embedding layer
        self.embedding = nn.Embedding(vocab_size, embedding_dim)

        # lstm layer
        self.lstm = nn.LSTM(embedding_dim,
                            hidden_dim,
                            num_layers = n_layers,
                            bidirectional=bidirectional,
                            dropout=dropout,
                            batch_first=True)
        # dense layer
        self.fc = nn.Linear(hidden_dim*2, output_dim)

        # activation
        self.act = nn.Sigmoid()

    def forward(self, text, text_lengths):
        # text=[batch size, sent_length]
        embedded = self.embedding(text)
        # embedded = [batch size, sent_len, emb dim]

        #packed sequence
        packed_embedded = nn.utils.rnn.pack_padded_sequence(embedded, 
                                                            text_lengths,
                                                            batch_first=True)
        packed_output, (hidden, cell) = self.lstm(packed_embedded)


        # concat the final forward and backward hidden state
        hidden = torch.cat((hidden[-2,:,:], hidden[-1,:,:]), dim=1)


        dense_outputs=self.fc(hidden)

        #final activation
        outputs = self.act(dense_outputs)
        return outputs

# Adopted from train.py
def train(model, iterator, optimizer, criterion):

    # initialize every epoch
    epoch_loss = 0
    epoch_acc = 0

    # set the model in training mode
    model.train()

    for batch in iterator:
        #reset the gradients after every batch
        optimizer.zero_grad()

        # retrieve text and num of words
        text, text_lengths = batch.text

        # convert to 1D tensor
        predictions = model(text, text_lengths).squeeze()

        # compute the loss
        loss = criterion(predictions, batch.label)
        # print(loss.item())
        # compute binary accuracy
        acc = binary_accuracy(predictions, batch.label)
        # print(acc.item())

        # backpropogate the loss and compute the gradients
        loss.backward()

        # update weights
        optimizer.step()

        # loss and accuracy
        epoch_loss += loss.item()
        epoch_acc += acc.item()
    return epoch_loss/len(iterator), epoch_acc/len(iterator)

def evaluate(model, iterator, criterion):

    # initialize every epoch
    epoch_loss = 0
    epoch_acc = 0

    #set model to eval mode
    model.eval()

    # deactivates auto_grad
    with torch.no_grad():
        for batch in iterator:
            # retrieve text and num of words
            text, text_lengths = batch.text 

            # convert to 1D tensor
            predictions = model(text, text_lengths).squeeze()

            loss = criterion(predictions, batch.label)
            acc = binary_accuracy(predictions, batch.label)

            # keep track of loss and accuracy
            epoch_loss += loss.item()
            epoch_acc += acc.item()

    return epoch_loss/len(iterator), epoch_acc/len(iterator)

# Adopted from train.py
optimizer = optim.Adam(model.parameters())
criterion = nn.BCELoss()
N_EPOCHS = config.N_EPOCH
best_valid_loss = float('inf')

# start training
for epoch in range(N_EPOCHS):

    # train the model
    train_loss, train_acc = train(model, train_iterator, optimizer, criterion)

    # evaluate the model
    valid_loss, valid_acc = evaluate(model, valid_iterator, criterion)

    # save the best model
    if valid_loss < best_valid_loss:
        best_valid_loss = valid_loss
        # torch.save(model.state_dict(), './output/best_model.pt')
        torch.save(model, './output/best_model.pt')
    print("Epoch {}/{}:".format(epoch+1, N_EPOCHS))
    print(f'\tTrain Loss: {train_loss:.3f} | Train Acc: {train_acc*100:.2f}%')
    print(f'\t Val. Loss: {valid_loss:.3f} |  Val. Acc: {valid_acc*100:.2f}%')