code fixes for new model

app.py

@@ -71,7 +71,7 @@ def get_words(doc_emb):
 
         #print(lang, end=": ")
 
-        scores = mul_model.E[lang] @ (doc_emb).T
+        scores = mul_model.E[lang].detach().numpy() @ (doc_emb).T
         k_ixs = np.argsort(scores)[::-1][:topn].squeeze()  # sort them in descending order and pick topn
         tmp = []
         for i in k_ixs:

learn_multi_doc_model.py

@@ -8,13 +8,15 @@ 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:
+class Model(torch.nn.Module):
     """Model defintion, parameters and helper fucntions to compute log-likelihood"""
 
     def __init__(self, vocab: dict, emb_dim: int):
@@ -25,23 +27,31 @@ class Model:
             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 = {}
+        # self.E = {} # torch.nn.ModuleDict
+        # self.E = torch.nn.ModuleDict()
+        self.E = torch.nn.ParameterDict()
 
         # bias vector for each language
-        self.b = {}
+        # 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] = np.random.uniform(-n2, n1, size=(vocab_size, emb_dim))
-            self.b[lang] = np.random.randn(vocab_size, 1) * 0.0001
+            # 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.
@@ -56,9 +66,13 @@ class Model:
         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)
+        # 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`
@@ -70,14 +84,17 @@ class Model:
             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] @ 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, axis=1)  # shape is n_docs x vocab_size
+        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]
@@ -114,9 +131,22 @@ class Model:
             llh = (X * log_thetas).sum()
         else:
             # X is a scipy sparse matrix
-            llh = (X.multiply(log_thetas)).sum()
+            # 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
+        return llh * (-1.0)  # * -1.0 when using pytorch to get negative llh (loss)
 
 
 def gradients_WE(model, lang, DE_lang, X, alpha):
@@ -156,9 +186,13 @@ def gradients_WE(model, lang, DE_lang, X, alpha):
     ef_grads = np.zeros_like(model.E)
 
     tmp = (
-        X - np.multiply(X.sum(axis=1).reshape(-1, 1), np.exp(log_thetas))
+            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()
+    # 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))
@@ -181,14 +215,14 @@ def update_parameters(params, gradient, learning_rate):
     """
 
     assert (
-        params.shape == gradient.shape
+            params.shape == gradient.shape
     ), "The params and gradient must have same shape, \
     ({:d}, {:d}) != ({:d} {:d})".format(
         *params.shape, *gradient.shape
     )
 
-    new_params = params + (
-        learning_rate * gradient
+    new_params = params.detach() + (
+            learning_rate * gradient
     )  # since we are doing gradient ascent
     return new_params
 
@@ -197,6 +231,7 @@ 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():
@@ -209,7 +244,13 @@ def train(model, bow, DE, args):
 
         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)
 
@@ -217,6 +258,13 @@ def train(model, bow, DE, args):
 
             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
@@ -229,18 +277,18 @@ def train(model, bow, DE, args):
                 "Instead it decreased, which means the updates have overshooted.",
                 "Halving the learning_rate.",
             )
-            learning_rate = learning_rate * 0.5
+            #learning_rate = learning_rate * 0.5
 
         llhs.append(llh_ei)
 
-        # learning_rate scheduler
+        # 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 % 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"
+                    os.path.join(args.out_dir, f"model_{args.alpha}_{i}.pkl"), "wb"
             ) as fpw:
                 pickle.dump(model, fpw)
             np.savetxt(
@@ -283,12 +331,13 @@ def main():
 
             # 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]
+                    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)
 
@@ -304,7 +353,7 @@ def main():
     model, llhs = train(model, bows, doc_embs, args)
 
     with open(
-        os.path.join(args.out_dir, f"model_{args.alpha}_{args.epochs}.pkl"), "wb"
+            os.path.join(args.out_dir, f"model_{args.alpha}_{args.epochs}.pkl"), "wb"
     ) as fpw:
         pickle.dump(model, fpw)
 
@@ -317,7 +366,6 @@ def main():
 
 
 def parse_arguments():
-
     parser = argparse.ArgumentParser(
         description=__doc__, formatter_class=argparse.ArgumentDefaultsHelpFormatter
     )

topic_discovery/{cvect_25000_ar.pkl → cvect_100000_ar.pkl}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:b37e9e016646662718993e2368f9e88c4c21141f8944f23449f27c6d59e03221
-size 3047285
+oid sha256:bc2a5a244e0554ffd18efee81ece7f4997136eb7caa6fdaa142e44e264983291
+size 3232315

topic_discovery/{cvect_25000_bn.pkl → cvect_100000_bn.pkl}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:05b3adf720d522a38762fda2bb6da2c948389a437b2138004698d326181d971d
-size 157149
+oid sha256:f27fa832a05cf8d6b2389a83e490526bf54763940b49f817f3ace830b332200c
+size 125341

topic_discovery/{cvect_25000_de.pkl → cvect_100000_de.pkl}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:e551d8934e6a8e23c841437805bbed1b0e17eb2f3ab3e260b9104c1e30f452ad
-size 2037400
+oid sha256:7ea2571476bac39c97563a7d2ec94879de43abd9877edbbadb76dcda600167f5
+size 1869324

topic_discovery/{cvect_25000_el.pkl → cvect_100000_el.pkl}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:d5419f509f5666ae55a7f5cdfb1cf7ea41f3fa102ec639c19c4aeea8b2dffe32
-size 3681045
+oid sha256:10fa08af8eccba5723694050feffdc7fdc7c698d7b62502b499534e7493d8ab1
+size 4068227

topic_discovery/cvect_100000_en.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:25da04442b2ff7ffae9bbe69f09e464aadf4a389c70fe6f283e9b0309d636a81
+size 5019023

topic_discovery/cvect_100000_es.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:70e1d60edf8e6b09736968209cf476fd60db5b12e3ad221593ac02061ec13307
+size 5396926

topic_discovery/cvect_100000_fr.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:34af4a4569ccd5e82c9218a288fe92091d3d7444f3b3a570707dd03ad464150c
+size 3513313

topic_discovery/cvect_100000_it.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2135701318838fcee53219515e369ef7b2cb9451884b18e734ee840372f34bd7
+size 2810588

topic_discovery/cvect_100000_jp.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8dbb2536a0f269c7ab9e71f38cffb69f3fd5925d7b43eabcdc5050f1bea5b6f9
+size 3040253

topic_discovery/cvect_100000_mk.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a3ef12ac4c0952ca01c9191cb78c2e63e80b77c8a1faaae05133ff23ad26f161
+size 1931126

topic_discovery/cvect_100000_nl.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:39d3baf77620a5b2a4c8084af2d442595a968cad90b05653fc328870171e3733
+size 1159719

topic_discovery/cvect_100000_pl.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:dcde383c79f519a4bce289379008b31fd35598df920cb6ad55fa3a2aa305a56d
+size 1981167

topic_discovery/cvect_100000_pt.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a363d8aa2f5b1ef3877d429e02bd7823374989d5f10ed9814c574860e317b698
+size 2068800

topic_discovery/cvect_100000_ru.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6beff3c5bd38755ac0ab8d34b6d86589d2777ed9361f0c59cdda9e7d04ef6031
+size 6251799

topic_discovery/cvect_100000_zhs.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d4d1f8f6d9b0d37d6a2c03df52711d4de319a74ef6b0d5229e2fe7f7115f3fe6
+size 9212102

topic_discovery/cvect_100000_zht.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:34538fb71bf08905b8461484e152c2f950515c03e8ffc0e427b1ec1db5ee3406
+size 9724371

topic_discovery/cvect_25000_en.pkl

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:cb0ee36e4ef6738d408e30132c5d970be2e05728c305fccce06dc67b3941bea2
-size 4143980

topic_discovery/cvect_25000_es.pkl

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:d28eb842e6f4717a791de9c8c61014131dbea8d26f84f90c62cd54b05595a1c9
-size 4235561

topic_discovery/cvect_25000_fr.pkl

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:74ff26b2269c2033f78ecb1e5870c449423d42d668975e5e98e899b6d2489f64
-size 2967490

topic_discovery/cvect_25000_it.pkl

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:4e8892d88fd88e0d9e121e57e1b77810e47d34909944b2e65e2094d426f17daa
-size 2477565

topic_discovery/cvect_25000_jp.pkl

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:2c075e83209a4a23afe290aef6a301717f4eadfd118a278114ea142fdf882c20
-size 3082086

topic_discovery/cvect_25000_mg.pkl

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:958dd98498097b8463b1fbc6f068b512650d40397b9e53659dc2238032126181
-size 3643714

topic_discovery/cvect_25000_mk.pkl

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:6758e48f3626b7c91b7359097d27aedb6beaeb36c6a6632901c3fae3f6da5ea3
-size 2152452

topic_discovery/cvect_25000_nl.pkl

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:5f81d4942757d07cde33715cd00fe150c377b19070f57cc992230b8c6eeacb06
-size 1466263

topic_discovery/cvect_25000_pl.pkl

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:ad1d1d8853aa424ba47c81d52ab6fdd708d1a440901652d680482d092a88a44a
-size 2063425

topic_discovery/cvect_25000_pt.pkl

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:baef6e3fe017ed4feb3ac2e08701b77b4425ade9f39d700ab3d1b4a2d89059d6
-size 2001188

topic_discovery/cvect_25000_ru.pkl

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:89bfa381364b0df772b0a181df8740bf597733e328410c464e6690d58e8e212f
-size 5482015

topic_discovery/cvect_25000_zhs.pkl

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:1369c082d071340da56006eef8ffc380625c39fef4a7034b7d1e2927b1f54717
-size 9390903

topic_discovery/cvect_25000_zht.pkl

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:030a1c4b66cfecf4645de14f77d90d56886e8927225581c94e45a93006c0c633
-size 9965443

topic_discovery/cvects.key

@@ -1,17 +1,17 @@
-en topic_discovery/cvect_25000_en.pkl
-es topic_discovery/cvect_25000_es.pkl
-fr topic_discovery/cvect_25000_fr.pkl
-mg topic_discovery/cvect_25000_mg.pkl
-it topic_discovery/cvect_25000_it.pkl
-el topic_discovery/cvect_25000_el.pkl
-zhs topic_discovery/cvect_25000_zhs.pkl
-zht topic_discovery/cvect_25000_zht.pkl
-bn topic_discovery/cvect_25000_bn.pkl
-ru topic_discovery/cvect_25000_ru.pkl
-pt topic_discovery/cvect_25000_pt.pkl
-ar topic_discovery/cvect_25000_ar.pkl
-de topic_discovery/cvect_25000_de.pkl
-jp topic_discovery/cvect_25000_jp.pkl
-mk topic_discovery/cvect_25000_mk.pkl
-pl topic_discovery/cvect_25000_pl.pkl
-nl topic_discovery/cvect_25000_nl.pkl
+en topic_discovery/cvect_100000_en.pkl
+es topic_discovery/cvect_100000_es.pkl
+fr topic_discovery/cvect_100000_fr.pkl
+mg topic_discovery/cvect_100000_mg.pkl
+it topic_discovery/cvect_100000_it.pkl
+el topic_discovery/cvect_100000_el.pkl
+zhs topic_discovery/cvect_100000_zhs.pkl
+zht topic_discovery/cvect_100000_zht.pkl
+bn topic_discovery/cvect_100000_bn.pkl
+ru topic_discovery/cvect_100000_ru.pkl
+pt topic_discovery/cvect_100000_pt.pkl
+ar topic_discovery/cvect_100000_ar.pkl
+de topic_discovery/cvect_100000_de.pkl
+jp topic_discovery/cvect_100000_jp.pkl
+mk topic_discovery/cvect_100000_mk.pkl
+pl topic_discovery/cvect_100000_pl.pkl
+nl topic_discovery/cvect_100000_nl.pkl