<gh_stars>1-10
# coding: utf8

import sys
sys.path.append("../../")

import os
import torch
from tqdm import tqdm_notebook
from torch import nn
from torch import optim
from argparse import Namespace
from examples.commons import helper
from examples.doc_classification_with_pretrained_embeddings.cnn import CNN
from examples.doc_classification_with_pretrained_embeddings.news_dataset import NewsDataset


args = Namespace(
    news_csv="../../data/ag_news/output_munged.csv",
    vectorizer_file="vectorizer.json",
    model_state_file="model.pth",
    save_dir="../../model_storage/doc_classification",
    glove_filepath="../../data/glove/glove.6B.300d.txt",
    use_glove=True,
    embedding_size=300,
    hidden_dim=100,
    num_channels=100,
    seed=7676,
    learning_rate=1e-3,
    dropout_p=0.1,
    batch_size=128,
    num_epochs=100,
    early_stopping_criteria=5,
    cuda=True,
    catch_keyboard_interrupt=True,
    reload_from_files=False,
    expand_filepaths_to_save_dir=True
)


def predict_category(title, classifier, vectorizer, max_length):
    title = helper.preprocess_text(title)
    vectorized_title = torch.tensor(vectorizer.vectorize(title, vector_length=max_length))
    result = classifier(vectorized_title.unsqueeze(0), apply_softmax=True)

    probability_values, indices = result.max(dim=1)
    predicted_category = vectorizer.category_vocab.lookup_index(indices.item())

    return {"category": predicted_category,
            "probability": probability_values.item()}


def get_samples(k=5):
    samples = {}
    for cat in dataset.val_df.category.unique():
        samples[cat] = dataset.val_df.title[dataset.val_df.category == cat].tolist()[:k]
    return samples


# Prep work
if args.expand_filepaths_to_save_dir:
    args.vectorizer_file = os.path.join(args.save_dir, args.vectorizer_file)
    args.model_state_file = os.path.join(args.save_dir, args.model_state_file)
    print("Expanded filepaths:")
    print("\t{}".format(args.vectorizer_file))
    print("\t{}".format(args.model_state_file))

# Check CUDA
if not torch.cuda.is_available():
    args.cuda = False
args.device = torch.device("cuda" if args.cuda else "cpu")
print("Using CUDA: {}".format(args.cuda))

# Set seed for reproducibility
helper.set_seed_everywhere(args.seed, args.cuda)

# handle dirs
helper.handle_dirs(args.save_dir)


# Initializations
if args.reload_from_files:
    # training from a checkpoint
    print("Loading dataset and vectorizer")
    dataset = NewsDataset.load_dataset_and_load_vectorizer(
        args.news_csv, args.vectorizer_file)
else:
    print("Loading dataset and creating vectorizer")
    # create dataset and vectorizer
    dataset = NewsDataset.load_dataset_and_make_vectorizer(args.news_csv)
    dataset.save_vectorizer(args.vectorizer_file)

vectorizer = dataset.get_vectorizer()


# Use GloVe or randomly initialized embeddings
if args.use_glove:
    words = vectorizer.title_vocab._token_to_idx.keys()
    embeddings = helper.make_embedding_matrix(glove_filepath=args.glove_filepath,
                                              words=words)
    print("Using pre-trained embeddings")
else:
    embeddings = None
    print("Not using pre-trained embeddings")


classifier = CNN(embedding_size=args.embedding_size,
                 num_embeddings=len(vectorizer.title_vocab),
                 num_channels=args.num_channels,
                 hidden_dim=args.hidden_dim,
                 num_classes=len(vectorizer.category_vocab),
                 dropout_p=args.dropout_p,
                 pretrained_embeddings=embeddings,
                 padding_idx=0)

classifier = classifier.to(args.device)
dataset.class_weights = dataset.class_weights.to(args.device)


# loss and optimier
loss_func = nn.CrossEntropyLoss()
optimizer = optim.Adam(classifier.parameters(), lr=args.learning_rate)
scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer=optimizer,
                                                 mode="min",
                                                 factor=0.5,
                                                 patience=1)


# Training loop
train_state = helper.make_train_state(args)

epoch_bar = tqdm_notebook(desc="training routine",
                          total=args.num_epochs,
                          position=0)

dataset.set_split("train")
train_bar = tqdm_notebook(desc="split=train",
                          total=dataset.get_num_batches(args.batch_size),
                          position=1,
                          leave=True)
dataset.set_split("val")
val_bar = tqdm_notebook(desc="split=val",
                        total=dataset.get_num_batches(args.batch_size),
                        position=1,
                        leave=True)

try:
    for epoch_index in range(args.num_epochs):
        train_state["epoch_index"] = epoch_index

        dataset.set_split("train")
        batch_generator = helper.generate_batches(dataset,
                                                  batch_size=args.batch_size,
                                                  device=args.device)
        running_loss = 0.
        running_acc = 0.
        classifier.train()

        for batch_index, batch_dict in enumerate(batch_generator):
            optimizer.zero_grad()

            y_pred = classifier(batch_dict["x_data"])

            loss = loss_func(y_pred, batch_dict["y_target"])
            loss_t = loss.item()
            running_loss += (loss_t - running_loss) / (batch_index + 1)

            loss.backward()

            optimizer.step()

            acc_t = helper.compute_accuracy(y_pred, batch_dict["y_target"])
            running_acc += (acc_t - running_acc) / (batch_index + 1)

            train_bar.set_postfix(loss=running_loss,
                                  acc=running_acc,
                                  epoch=epoch_index)
            train_bar.update()

        train_state["train_loss"].append(running_loss)
        train_state["train_acc"].append(running_acc)

        dataset.set_split("val")
        batch_generator = helper.generate_batches(dataset,
                                                  batch_size=args.batch_size,
                                                  device=args.device)
        running_loss = 0.
        running_acc = 0.
        classifier.eval()

        for batch_index, batch_dict in enumerate(batch_generator):
            y_pred = classifier(batch_dict["x_data"])

            loss = loss_func(y_pred, batch_dict["y_target"])
            loss_t = loss.item()
            running_loss += (loss_t - running_loss) / (batch_index + 1)

            acc_t = helper.compute_accuracy(y_pred, batch_dict["y_target"])
            running_acc += (acc_t - running_acc) / (batch_index + 1)

            val_bar.set_postfix(loss=running_loss,
                                acc=running_acc,
                                epoch=epoch_index)
            val_bar.update()

        train_state["val_loss"].append(running_loss)
        train_state["val_acc"].append(running_acc)

        train_state = helper.update_train_state(args=args,
                                                model=classifier,
                                                train_state=train_state)

        scheduler.step(train_state["val_loss"][-1])

#         if train_state["stop_early"]:
#             break

        train_bar.n = 0
        val_bar.n = 0
        epoch_bar.update()

        print(f"Epoch {epoch_index}")
except KeyboardInterrupt:
    print("Exiting loop")


classifier.load_state_dict(torch.load(train_state["model_filename"]))

classifier = classifier.to(args.device)
dataset.class_weights = dataset.class_weights.to(args.device)
loss_func = nn.CrossEntropyLoss(dataset.class_weights)

dataset.set_split("test")
batch_generator = helper.generate_batches(dataset,
                                          batch_size=args.batch_size,
                                          device=args.device)

running_loss = 0.
running_acc = 0.
classifier.eval()

for batch_index, batch_dict in enumerate(batch_generator):
    y_pred = classifier(batch_dict["x_data"])

    loss = loss_func(y_pred, batch_dict["y_target"])
    loss_t = loss.item()
    running_loss += (loss_t - running_loss) / (batch_index + 1)

    acc_t = helper.compute_accuracy(y_pred, batch_dict["y_target"])
    running_acc += (acc_t - running_acc) / (batch_index + 1)

train_state["test_loss"] = running_loss
train_state["test_acc"] = running_acc

print(f"Test loss: {train_state['test_loss']}")
print(f"Test acc: {train_state['test_acc']}")


classifier = classifier.to("cpu")
val_samples = get_samples()

for truth, sample_group in val_samples.items():
    print(f"True Category: {truth}")
    print("=" * 30)
    for sample in sample_group:
        prediction = predict_category(sample,
                                      classifier,
                                      vectorizer,
                                      dataset._max_seq_length + 1)
        print("Prediction: {} (p={:.2f})".format(prediction["category"],
                                                 prediction["probability"]))
        print("\t + Sample: {}".format(sample))
    print("-" * 30)

# coding: utf8

import sys
sys.path.append("../../")

import os
import torch
from tqdm import tqdm_notebook
from torch import nn
from torch import optim
from argparse import Namespace
from examples.commons import helper
from examples.doc_classification_with_pretrained_embeddings.cnn import CNN
from examples.doc_classification_with_pretrained_embeddings.news_dataset import NewsDataset


args = Namespace(
    news_csv="../../data/ag_news/output_munged.csv",
    vectorizer_file="vectorizer.json",
    model_state_file="model.pth",
    save_dir="../../model_storage/doc_classification",
    glove_filepath="../../data/glove/glove.6B.300d.txt",
    use_glove=True,
    embedding_size=300,
    hidden_dim=100,
    num_channels=100,
    seed=7676,
    learning_rate=1e-3,
    dropout_p=0.1,
    batch_size=128,
    num_epochs=100,
    early_stopping_criteria=5,
    cuda=True,
    catch_keyboard_interrupt=True,
    reload_from_files=False,
    expand_filepaths_to_save_dir=True
)


def predict_category(title, classifier, vectorizer, max_length):
    title = helper.preprocess_text(title)
    vectorized_title = torch.tensor(vectorizer.vectorize(title, vector_length=max_length))
    result = classifier(vectorized_title.unsqueeze(0), apply_softmax=True)

    probability_values, indices = result.max(dim=1)
    predicted_category = vectorizer.category_vocab.lookup_index(indices.item())

    return {"category": predicted_category,
            "probability": probability_values.item()}


def get_samples(k=5):
    samples = {}
    for cat in dataset.val_df.category.unique():
        samples[cat] = dataset.val_df.title[dataset.val_df.category == cat].tolist()[:k]
    return samples


# Prep work
if args.expand_filepaths_to_save_dir:
    args.vectorizer_file = os.path.join(args.save_dir, args.vectorizer_file)
    args.model_state_file = os.path.join(args.save_dir, args.model_state_file)
    print("Expanded filepaths:")
    print("\t{}".format(args.vectorizer_file))
    print("\t{}".format(args.model_state_file))

# Check CUDA
if not torch.cuda.is_available():
    args.cuda = False
args.device = torch.device("cuda" if args.cuda else "cpu")
print("Using CUDA: {}".format(args.cuda))

# Set seed for reproducibility
helper.set_seed_everywhere(args.seed, args.cuda)

# handle dirs
helper.handle_dirs(args.save_dir)


# Initializations
if args.reload_from_files:
    # training from a checkpoint
    print("Loading dataset and vectorizer")
    dataset = NewsDataset.load_dataset_and_load_vectorizer(
        args.news_csv, args.vectorizer_file)
else:
    print("Loading dataset and creating vectorizer")
    # create dataset and vectorizer
    dataset = NewsDataset.load_dataset_and_make_vectorizer(args.news_csv)
    dataset.save_vectorizer(args.vectorizer_file)

vectorizer = dataset.get_vectorizer()


# Use GloVe or randomly initialized embeddings
if args.use_glove:
    words = vectorizer.title_vocab._token_to_idx.keys()
    embeddings = helper.make_embedding_matrix(glove_filepath=args.glove_filepath,
                                              words=words)
    print("Using pre-trained embeddings")
else:
    embeddings = None
    print("Not using pre-trained embeddings")


classifier = CNN(embedding_size=args.embedding_size,
                 num_embeddings=len(vectorizer.title_vocab),
                 num_channels=args.num_channels,
                 hidden_dim=args.hidden_dim,
                 num_classes=len(vectorizer.category_vocab),
                 dropout_p=args.dropout_p,
                 pretrained_embeddings=embeddings,
                 padding_idx=0)

classifier = classifier.to(args.device)
dataset.class_weights = dataset.class_weights.to(args.device)


# loss and optimier
loss_func = nn.CrossEntropyLoss()
optimizer = optim.Adam(classifier.parameters(), lr=args.learning_rate)
scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer=optimizer,
                                                 mode="min",
                                                 factor=0.5,
                                                 patience=1)


# Training loop
train_state = helper.make_train_state(args)

epoch_bar = tqdm_notebook(desc="training routine",
                          total=args.num_epochs,
                          position=0)

dataset.set_split("train")
train_bar = tqdm_notebook(desc="split=train",
                          total=dataset.get_num_batches(args.batch_size),
                          position=1,
                          leave=True)
dataset.set_split("val")
val_bar = tqdm_notebook(desc="split=val",
                        total=dataset.get_num_batches(args.batch_size),
                        position=1,
                        leave=True)

try:
    for epoch_index in range(args.num_epochs):
        train_state["epoch_index"] = epoch_index

        dataset.set_split("train")
        batch_generator = helper.generate_batches(dataset,
                                                  batch_size=args.batch_size,
                                                  device=args.device)
        running_loss = 0.
        running_acc = 0.
        classifier.train()

        for batch_index, batch_dict in enumerate(batch_generator):
            optimizer.zero_grad()

            y_pred = classifier(batch_dict["x_data"])

            loss = loss_func(y_pred, batch_dict["y_target"])
            loss_t = loss.item()
            running_loss += (loss_t - running_loss) / (batch_index + 1)

            loss.backward()

            optimizer.step()

            acc_t = helper.compute_accuracy(y_pred, batch_dict["y_target"])
            running_acc += (acc_t - running_acc) / (batch_index + 1)

            train_bar.set_postfix(loss=running_loss,
                                  acc=running_acc,
                                  epoch=epoch_index)
            train_bar.update()

        train_state["train_loss"].append(running_loss)
        train_state["train_acc"].append(running_acc)

        dataset.set_split("val")
        batch_generator = helper.generate_batches(dataset,
                                                  batch_size=args.batch_size,
                                                  device=args.device)
        running_loss = 0.
        running_acc = 0.
        classifier.eval()

        for batch_index, batch_dict in enumerate(batch_generator):
            y_pred = classifier(batch_dict["x_data"])

            loss = loss_func(y_pred, batch_dict["y_target"])
            loss_t = loss.item()
            running_loss += (loss_t - running_loss) / (batch_index + 1)

            acc_t = helper.compute_accuracy(y_pred, batch_dict["y_target"])
            running_acc += (acc_t - running_acc) / (batch_index + 1)

            val_bar.set_postfix(loss=running_loss,
                                acc=running_acc,
                                epoch=epoch_index)
            val_bar.update()

        train_state["val_loss"].append(running_loss)
        train_state["val_acc"].append(running_acc)

        train_state = helper.update_train_state(args=args,
                                                model=classifier,
                                                train_state=train_state)

        scheduler.step(train_state["val_loss"][-1])

#         if train_state["stop_early"]:
#             break

        train_bar.n = 0
        val_bar.n = 0
        epoch_bar.update()

        print(f"Epoch {epoch_index}")
except KeyboardInterrupt:
    print("Exiting loop")


classifier.load_state_dict(torch.load(train_state["model_filename"]))

classifier = classifier.to(args.device)
dataset.class_weights = dataset.class_weights.to(args.device)
loss_func = nn.CrossEntropyLoss(dataset.class_weights)

dataset.set_split("test")
batch_generator = helper.generate_batches(dataset,
                                          batch_size=args.batch_size,
                                          device=args.device)

running_loss = 0.
running_acc = 0.
classifier.eval()

for batch_index, batch_dict in enumerate(batch_generator):
    y_pred = classifier(batch_dict["x_data"])

    loss = loss_func(y_pred, batch_dict["y_target"])
    loss_t = loss.item()
    running_loss += (loss_t - running_loss) / (batch_index + 1)

    acc_t = helper.compute_accuracy(y_pred, batch_dict["y_target"])
    running_acc += (acc_t - running_acc) / (batch_index + 1)

train_state["test_loss"] = running_loss
train_state["test_acc"] = running_acc

print(f"Test loss: {train_state['test_loss']}")
print(f"Test acc: {train_state['test_acc']}")


classifier = classifier.to("cpu")
val_samples = get_samples()

for truth, sample_group in val_samples.items():
    print(f"True Category: {truth}")
    print("=" * 30)
    for sample in sample_group:
        prediction = predict_category(sample,
                                      classifier,
                                      vectorizer,
                                      dataset._max_seq_length + 1)
        print("Prediction: {} (p={:.2f})".format(prediction["category"],
                                                 prediction["probability"]))
        print("\t + Sample: {}".format(sample))
    print("-" * 30)
# coding: utf8 # Prep work # Check CUDA # Set seed for reproducibility # handle dirs # Initializations # training from a checkpoint # create dataset and vectorizer # Use GloVe or randomly initialized embeddings # loss and optimier # Training loop #         if train_state["stop_early"]: #             break