learn_multi_doc_model.py

# /usr/bin/env python3

import argparse
import os
import numpy as np
import scipy
import pickle
from scipy.special import log_softmax
from time import time
from packaging import version
import torch

assert version.parse(scipy.__version__) >= version.parse(
    "1.7.0"
), f"Requries scipy > 1.7.0. Found {scipy.__version__}"

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

class Model(torch.nn.Module):
    """Model defintion, parameters and helper fucntions to compute log-likelihood"""

    def __init__(self, vocab: dict, emb_dim: int):
        """Initialize our model

        Args:
            vocab: vocab size for each language {'en': 25000, 'de': 25000}
            emb_dim: embedding dimension, will be same across languages
        """

        super().__init__()

        self.L = len(vocab)
        self.vocab = vocab
        self.emb_dim = emb_dim

        # word embeddings matrix / subspace for each language
        # self.E = {} # torch.nn.ModuleDict
        # self.E = torch.nn.ModuleDict()
        self.E = torch.nn.ParameterDict()

        # bias vector for each language
        # self.b = {} # torch.nn.ModuleDict
        # self.b = torch.nn.ModuleDict()
        self.b = torch.nn.ParameterDict()

        n1 = 1.0 / np.sqrt(emb_dim)

        # initialize word embeddings and bias vectors randomly
        for lang, vocab_size in vocab.items():
            n2 = 1.0 / np.sqrt(vocab_size)
            # self.E[lang] = torch.nn.ParameterList(torch.from_numpy(np.random.uniform(-n2, n1, size=(vocab_size, emb_dim))))
            self.E[lang] = torch.nn.Parameter(torch.Tensor(np.random.uniform(-n2, n1, size=(vocab_size, emb_dim))),
                                              requires_grad=True).to(device)
            self.b[lang] = torch.nn.Parameter(torch.Tensor(np.random.randn(vocab_size, 1) * 0.0001), requires_grad=True).to(device)

    def init_bias_with_log_unigram_dist(self, X, lang):
        """We will initialize the bias vector with log of unigram distribution over vocabulary.
        This should help us with better initialization.

        b = \log (\sum_d x_d) / (\sum_d \sum_i x_{di})
        """

        # if X is sparse matrix, X.A gives the dense version of it in numpy array format
        if isinstance(X, np.ndarray):
            X = X + 1e-08  # to avoid zeros
        else:
            X = X.A + 1e-08  # to avoid any zeros

        # self.b[lang][:, 0] = np.log(
        #    X.sum(axis=0) / X.sum()
        # )  # we would like b to of size (W, 1)

        b_copy = self.b[lang].clone()
        b_copy[:, 0] = torch.from_numpy(np.log(X.sum(axis=0) / X.sum()))
        self.b[lang] = torch.nn.Parameter(b_copy, requires_grad=True)

    def compute_log_thetas(self, lang: str, DE_lang: np.ndarray, sanity_check=False):
        """Compute log of thetas, where theta_d is the unigram distribution over document `d`
        estiamted from the current params (word-embedding matrix, bias vector) and document embedding a_d.

        Args:
        ----
            lang (str): Language ID (eg: en, de, es ...)
            DE_lang (np.ndarray): Document embeddings of language
        """

        # mat = self.b[lang] + (self.E[lang] @ DE_lang)  # shape is vocab_size x n_docs
        mat = self.b[lang] + (self.E[lang].double() @ torch.from_numpy(DE_lang).double().to(device))
        # mat = mat.detach()
        # mat = mat.detach().T
        mat = mat.T  # shape is D x W

        # log_norm = logsumexp(mat, axis=1)
        # log_thetas = mat - log_norm

        # the following single step is same the two above steps combined
        log_thetas = log_softmax(mat.detach().numpy(), axis=1)  # shape is n_docs x vocab_size

        if sanity_check:
            n_docs = DE_lang.shape[0]
            # sanity-check
            # since each document is a proper distribution, it should sum upto 1
            # sum of the matrix should be equal to number of documents
            print(
                "Sanity check for log-thetas:",
                np.allclose(np.exp(log_thetas).sum(), n_docs),
            )

        return log_thetas

    def compute_log_likelihood(self, lang, DE_lang, X):
        """Compute log-likelihood of the data, given the current parameters / embeddings

        Each summation could be implemented using a for-loop but that would very slow,
        since we have every thing stored in matrices and a sparse matrix, we will do it via
        matrix muliplications and additions.

        Args:
            lang: language ID (eg: en, es, fr)
            DE_lang: document embeddings for the given language
            X: doc-by-word counts in scipy.sparse format for a specific language

        Returns:
            float: log-likelihood of the data
        """

        log_thetas = self.compute_log_thetas(lang, DE_lang)

        # log-likelihood is product of counts to the respective log-probability values.
        if isinstance(X, np.ndarray):
            llh = (X * log_thetas).sum()
        else:
            # X is a scipy sparse matrix
            # this is the tricky part in pytorch

            coo = X.tocoo()

            row_ixs = torch.LongTensor(coo.row).to(device)
            col_ixs = torch.LongTensor(coo.col).to(device)
            data = torch.FloatTensor(coo.data).to(device)

            # llh = (X.multiply(log_thetas)).sum()

            log_thetas_tensor = torch.from_numpy(log_thetas)

            llh = (log_thetas_tensor[row_ixs, col_ixs] * data).sum()
            # TODO row_ixs, col_ixs, data

        return llh * (-1.0)  # * -1.0 when using pytorch to get negative llh (loss)


def gradients_WE(model, lang, DE_lang, X, alpha):
    """Gradient of the log-likelihood with-respect-to language-specific word embedding matrix `E`

    Args:
        model (Model): The object of the model
        lang (str): Language ID
        DE_lang: document embeddings for the given language
        X (scipy.sparse_matrix): The doc-by-word counts
        alpha (float): L2 reg. weight

    Returns:
        np.ndarray: Gradient of log-likelihood w.r.t word embeddings, i.e, grad of llh w.r.t to model.E
    """

    # grads = np.zeros_like(model.E)  # initialize empty gradients to be the same shape as word embeddings (W, K)

    # compute log_thetas as they are needed in gradient
    log_thetas = model.compute_log_thetas(lang, DE_lang)

    # the gradient computation can be done using for-loops to reflect the equation
    # or it can be done efficiently using matrix multiplications

    # 1. simple way using for-loop
    # iterate over all documents
    # for d in range(model.D):

    # iterate over every word,
    #     for k in range(model.W):
    #         x_dk = X[d, k]  # count of word k in doc d
    #         rel_x_dk = X[d, :].sum() * np.exp(log_thetas)[d, k]  # relative /estimated count of word k in doc d
    #         grads[k, :] += ((x_dk - rel_x_dk) * model.A[:, d])  # doc embeddings are column wise in model.A

    # 2. Efficient way of obtaining gradients using matrix operations

    ef_grads = np.zeros_like(model.E)

    tmp = (
            X - np.multiply(X.sum(axis=1).reshape(-1, 1), np.exp(log_thetas))
    ).A  # .A will convert matrix to np ndarray
    # ef_grads = (DE_lang @ tmp).T - (alpha * 0.5 * model.E[lang]).sum(axis=1, keepdims=True)

    m = model.E[lang].detach().numpy()
    # ef_grads = (DE_lang @ tmp).T - (alpha * 0.5 * model.E[lang]).sum(axis=1, keepdims=True)
    ef_grads = (DE_lang @ tmp).T - (alpha * 0.5 * m).sum(axis=1, keepdims=True)

    # Sanity check to see if gradients computed in both ways are numerically identical
    # print('- All close grad_E:', np.allclose(ef_grads, grads))

    return ef_grads


def update_parameters(params, gradient, learning_rate):
    """Update the parameters

    Args:
        params (np.ndarray): Word embedding matrix of the document embedding matrix
        gradient (np.ndarray): Gradients of all word embeddings or document embeddings. Should be same as size as params
        learning_rate (float): The learning_rate can also be seen as step size, i.e, the size of the step to be taken
               along the direction of gradient. Too big steps can overshoot our estimate, whereas too small steps
               can take longer for the model to reach optimum.

    Returns:
        np.ndarray: the updated params
    """

    assert (
            params.shape == gradient.shape
    ), "The params and gradient must have same shape, \
    ({:d}, {:d}) != ({:d} {:d})".format(
        *params.shape, *gradient.shape
    )

    new_params = params.detach() + (
            learning_rate * gradient
    )  # since we are doing gradient ascent
    return new_params


def train(model, bow, DE, args):
    """Training scheme for the model"""

    print("\nTraining started ..")
    optim = torch.optim.Adam(model.parameters(), lr=args.lr)
    learning_rate = args.lr
    llh_0 = 0.0
    for lang, X in bow.items():
        llh_0 += model.compute_log_likelihood(lang, DE[lang].T, X)
    print("  Initial log-likelihood: {:16.2f}".format(llh_0))

    llhs = [llh_0]

    for i in range(1, args.epochs + 1):

        llh_ei = 0.0
        for lang, X in bow.items():
            # for pytorch
            optim.zero_grad()
            # get row_ixs, col_ixs, data from X

            # compute neg llh
           #loss = torch.tensor(llh_ei, requires_grad=True)
            #loss = torch.as_tensor(llh_ei).detach().clone()
            # update word embeddings E for lang, by keeping doc-embeddings A fixed
            grad_E = gradients_WE(model, lang, DE[lang].T, X, args.alpha)

            model.E[lang] = update_parameters(model.E[lang], grad_E, learning_rate)

            llh_ei += model.compute_log_likelihood(lang, DE[lang].T, X)

            loss = torch.tensor(llh_ei, requires_grad=True)
            loss.backward()

            optim.step()



        print(
            "Epoch {:4d} / {:4d} | Log-likelihood: {:16.2f} | Learning rate: {:f}".format(
                i, args.epochs, llh_ei, learning_rate
            )
        )

        if llh_ei < llhs[-1]:
            print(
                "The log-likelihood should improve after every epoch.",
                "Instead it decreased, which means the updates have overshooted.",
                "Halving the learning_rate.",
            )
            #learning_rate = learning_rate * 0.5

        llhs.append(llh_ei)

        # ylearning_rate scheduler
        # we reduce the learning_rate by 10 % after every 10 epochs
        if i % 10 == 0:
           print("Reducing the learning by a factor of 0.1 every 10 epcohs")
           learning_rate -= learning_rate * 0.1
        if i % 100 == 0:
            with open(
                    os.path.join(args.out_dir, f"model_{args.alpha}_{i}.pkl"), "wb"
            ) as fpw:
                pickle.dump(model, fpw)
            np.savetxt(
                os.path.join(args.out_dir, f"llh_{args.alpha}_{args.epochs}.txt"),
                np.asarray(llhs),
            )

    return model, llhs


def main():
    """main"""

    args = parse_arguments()

    os.makedirs(args.out_dir, exist_ok=True)

    emb_dim = 0
    # load doc embeddings for each language
    doc_embs = {}  # {lang_1: np.ndarray, lang_2: np.ndarray, ...}
    with open(args.input_embedding_key_file, "r") as fpr:
        for line in fpr:
            lang, fpath = line.strip().split()
            doc_embs[lang] = np.load(fpath)
            print("Loaded embeddings:", lang, doc_embs[lang].shape)

            if emb_dim == 0:
                emb_dim = doc_embs[lang].shape[1]

    # load bag of words for each language
    bows = {}  # {lang_1: scipy.sparse, lang_2: scipy.sparse, ...}
    vocab = {}  # {lang_1: vocab_size}
    with open(args.input_bag_of_words_key_file, "r") as fpr:
        for line in fpr:
            lang, fpath = line.strip().split()
            bows[lang] = scipy.sparse.load_npz(fpath)
            print("Loaded bag-of-words:", lang, bows[lang].shape)

            vocab[lang] = bows[lang].shape[1]

            # assert the number of docs per language are same in embeddings and bag-of-words
            assert (
                    bows[lang].shape[0] == doc_embs[lang].shape[0]
            ), "Number of docs in BoW ({:d}) != number of docs in embeddigs ({:d}) for language: {:s}".format(
                bows[lang].shape[0], doc_embs[lang].shape[0], lang
            )

    model = Model(vocab, emb_dim)
    model.to(device)
    for lang, bow in bows.items():
        model.init_bias_with_log_unigram_dist(bow, lang)

    print("Model params:")
    for lang in model.vocab:
        print("  ", lang, model.E[lang].shape, model.b[lang].shape)

    if args.resume:
        with open(args.resume, "rb") as fpr:
            model = pickle.load(fpr)

    # start the training
    model, llhs = train(model, bows, doc_embs, args)

    with open(
            os.path.join(args.out_dir, f"model_{args.alpha}_{args.epochs}.pkl"), "wb"
    ) as fpw:
        pickle.dump(model, fpw)

    np.savetxt(
        os.path.join(args.out_dir, f"llh_{args.alpha}_{args.epochs}.txt"),
        np.asarray(llhs),
    )

    print("Saved in", args.out_dir)


def parse_arguments():
    parser = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )

    parser.add_argument(
        "input_embedding_key_file",
        help="path to file that has paths to embeddings for each language",
    )

    parser.add_argument(
        "input_bag_of_words_key_file", help="path to input bag of words dictionary file"
    )

    parser.add_argument("out_dir", help="out dir to save the model/word embeddings")

    parser.add_argument("--epochs", type=int, default=100, help="number of epochs")
    parser.add_argument("--lr", type=float, default=0.0001, help="learning rate")
    parser.add_argument(
        "--alpha", type=float, default=1e-4, help="L2 reg. weight / weight decay"
    )

    parser.add_argument(
        "--resume", default="", help="path to trained model to resume training"
    )

    args = parser.parse_args()

    return args


if __name__ == "__main__":
    main()