Training function that orchestrates the Classification!

def train(args):
    """Training function that orchestrates the Classification! """
    train_file = args.input
    test_file = args.validation
    ngram_range = args.ngrams
    logging.info('Ngrams range for the training run : %s', ngram_range)
    logging.info('Loading Training data')
    train_labels, train_data = read_input_data(train_file)
    logging.info('Loading Test data')
    test_labels, test_data = read_input_data(test_file)
    tokens_list = []
    for line in train_data:
        tokens_list.extend(line.split())

    cntr = Counter(tokens_list)
    train_vocab = gluonnlp.Vocab(cntr)
    logging.info('Vocabulary size: %s', len(train_vocab))
    logging.info('Training data converting to sequences...')
    embedding_matrix_len = len(train_vocab)
    # Preprocess the dataset
    train_sequences = convert_to_sequences(train_data, train_vocab)
    test_sequences = convert_to_sequences(test_data, train_vocab)

    if ngram_range >= 2:
        logging.info('Adding %s-gram features', ngram_range)
        # Create set of unique n-gram from the training set.
        ngram_set = set()
        for input_list in train_sequences:
            for i in range(2, ngram_range + 1):
                set_of_ngram = create_ngram_set(input_list, ngram_value=i)
                ngram_set.update(set_of_ngram)
        start_index = len(cntr)
        token_indices = {v: k + start_index for k, v in enumerate(ngram_set)}
        embedding_matrix_len = embedding_matrix_len + len(token_indices)
        train_sequences = add_ngram(train_sequences, token_indices,
                                    ngram_range)
        test_sequences = add_ngram(test_sequences, token_indices, ngram_range)
        logging.info('Added n-gram features to train and test datasets!! ')
    logging.info('Encoding labels')

    label_mapping = get_label_mapping(train_labels)
    y_train_final = list(map(lambda x: label_mapping[x], train_labels))
    y_test_final = list(map(lambda x: label_mapping[x], test_labels))

    train_sequences, train_data_lengths = preprocess_dataset(
        train_sequences, y_train_final)
    test_sequences, _ = preprocess_dataset(
        test_sequences, y_test_final)
    train_dataloader, test_dataloader = get_dataloader(train_sequences, train_data_lengths,
                                                       test_sequences, args.batch_size)

    num_classes = len(np.unique(train_labels))
    logging.info('Number of labels: %s', num_classes)
    logging.info('Initializing network')
    ctx = get_context(args)
    logging.info('Running Training on ctx:%s', ctx)
    embedding_dim = args.emsize
    logging.info('Embedding Matrix Length:%s', embedding_matrix_len)
    net = FastTextClassificationModel(
        embedding_matrix_len, embedding_dim, num_classes)
    net.hybridize()
    net.collect_params().initialize(mx.init.Xavier(), ctx=ctx)
    logging.info('Network initialized')

    softmax_cross_entropy = gluon.loss.SoftmaxCrossEntropyLoss()
    sigmoid_loss_fn = gluon.loss.SigmoidBinaryCrossEntropyLoss()

    loss_function = softmax_cross_entropy
    if num_classes == 2:
        logging.info(
            'Changing the loss function to sigmoid since its Binary Classification'
        )
        loss_function = sigmoid_loss_fn
    logging.info('Loss function for training:%s', loss_function)
    num_epochs = args.epochs
    batch_size = args.batch_size
    logging.info('Starting Training!')
    learning_rate = args.lr
    trainer = gluon.Trainer(net.collect_params(), 'adam',
                            {'learning_rate': learning_rate})
    num_batches = len(train_data) / batch_size
    display_batch_cadence = int(math.ceil(num_batches / 10))
    logging.info('Training on %s samples and testing on %s samples',
                 len(train_data), len(test_data))
    logging.info('Number of batches for each epoch : %s, Display cadence: %s',
                 num_batches, display_batch_cadence)
    for epoch in range(num_epochs):
        for batch, ((data, length), label) in enumerate(train_dataloader):
            data = data.as_in_context(ctx)
            label = label.as_in_context(ctx)
            length = length.astype('float32').as_in_context(ctx)
            with autograd.record():
                output = net(data, length)
                loss = loss_function(output, label)
            loss.backward()
            trainer.step(data.shape[0])
            if batch % display_batch_cadence == 0:
                logging.info('Epoch : %s, Batches complete :%s', epoch, batch)
        logging.info('Epoch complete :%s, Computing Accuracy', epoch)
        test_accuracy, test_loss = evaluate_accuracy(
            test_dataloader, net, ctx, loss_function, num_classes)
        logging.info('Epochs completed : %s Test Accuracy: %s, Test Loss: %s',
                     epoch, test_accuracy, test_loss)
        learning_rate = learning_rate * 0.5
        trainer.set_learning_rate(learning_rate)
    save_model(net, args.output)