# /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()