import numpy as np import pandas as pd import torch import torch.nn as nn from torch.nn.init import xavier_uniform_ from torch_model_base import TorchModelBase import utils __author__ = "Christopher Potts" __version__ = "CS224u, Stanford, Spring 2022" class TorchGloVeDataset(torch.utils.data.Dataset): def __init__(self, X_log, weights): """ Dataset class for GloVe. A separate class is needed only because, for batching, we need the batch indices for the `forward` method in `TorchGloVeModel`. For details on the construction of the two tensor arguments, see `TorchGloVe.fit`. Parameters ---------- X_log : torch.FloatTensor, shape `(n_vocab, n_vocab)` weights : torch.FloatTensor, shape `(n_vocab, n_vocab)` """ self.X_log = X_log self.weights = weights assert len(self.X_log) == len(self.weights) def __len__(self): return len(self.X_log) def __getitem__(self, idx): return self.X_log[idx], idx, self.weights[idx] class TorchGloVeModel(nn.Module): def __init__(self, n_words, embed_dim): super().__init__() self.n_words = n_words self.embed_dim = embed_dim self.W = self._init_weights(self.n_words, self.embed_dim) self.C = self._init_weights(self.n_words, self.embed_dim) self.bw = self._init_weights(self.n_words, 1) self.bc = self._init_weights(self.n_words, 1) def _init_weights(self, m, n): return nn.Parameter( xavier_uniform_(torch.empty(m, n))) def forward(self, X_log, idx): """ Parameters ---------- X_log : torch.FloatTensor, shape `(batch_size, n_vocab)`. idx : torch.LongTensor, shape `(batch_size, )` Indices of the vocab items in the current batch. Returns ------- torch.FloatTensor, shape `(n_vocab, n_vocab)`. """ preds = self.W[idx].matmul(self.C.T) + self.bw[idx] + self.bc.T diffs = preds - X_log return diffs class TorchGloVeLoss(nn.Module): def __init__(self): super().__init__() self.reduction = "sum" def forward(self, diffs, weights): return torch.sum(0.5 * torch.mul(weights, diffs**2)) class TorchGloVe(TorchModelBase): def __init__(self, embed_dim=100, alpha=0.75, xmax=100, **model_kwargs): """ Defines the GloVe computation graph using the vectorized method given in Dingwall, Nicholas and Christopher Potts. 2018. Mittens: An extension of GloVe for learning domain-specialized representations. Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, 212-217. """ super().__init__(**model_kwargs) self.embed_dim = embed_dim self.alpha = alpha self.xmax = xmax # We can't actually split the data into train and test portions, # but we can evaluate how we're doing holistically to try to # find the best model: if model_kwargs.get("early_stopping"): self.validation_fraction = 1.0 self.loss = TorchGloVeLoss() self.params += ['embed_dim', 'alpha', 'xmax'] def build_dataset(self, X_log, weights): """ Creates a simple `TorchGloVeDataset`, which is really just a TensorDataset that returns the examples indices as well as the core batch dataset structrues, so that we can do batch updates to the GloVe parameters. Parameters ---------- X_log : torch.FloatTensor, shape `(n_vocab, n_vocab)`. The log of the count matrix, ignoring 0s. See `fit. weights : torch.FloatTensor, shape `(n_vocab, n_vocab)` Returns ------- TorchGloVeDataset """ X_log = torch.FloatTensor(X_log) weights = torch.FloatTensor(weights) dataset = TorchGloVeDataset(X_log, weights) return dataset def build_graph(self): """ The core computation graph. Called by `fit` to set the attribute `model`. Returns ------- TorchGloVeModel """ return TorchGloVeModel(self.n_words, self.embed_dim) def fit(self, X): """ Prepares `X` to permit learning against the GloVe objective, and then uses the superclass `fit` method to train the model parameters. Unlike the supervised models in this repository, this method returns the learned embedding (W + C) rather than `self`, so that it acts like a model that transforms a vector space (see also the autoencoder models). Parameters ---------- X : np.array, shape `(n_words, n_words)` This should be a square matrix of possible scaled co-occurrence counts. Attributes ---------- self.embedding: np.array, shape (n_words, embed_dim) The same matrix that is returned by the method. Returns ------- embedding: np.array, shape (n_words, embed_dim) The same matrix that is stored as `self.embedding`. """ X_vals = self.convert_input_to_array(X) self.n_words = len(X_vals) # This applies the function # # f(x) = (x/self.xmax)**self.alpha if x < self.xmax, else 1.0 # # to the full count matrix: bounded = np.minimum(X_vals, self.xmax) weights = (bounded / self.xmax)**self.alpha # Precompute log X[i, j] for all i, j: X_log = utils.log_of_array_ignoring_zeros(X_vals) super().fit(X_log, weights) # Per the advice in the paper, use the sum of the word and # context embeddings: embedding = self.model.W + self.model.C embedding = embedding.detach().cpu().numpy() # If the input was a `pd.DataFrame`, return one as well: self.embedding = self.convert_output(embedding, X) return self.embedding def score(self, X, y=None): """ The goal of GloVe is to learn vectors whose dot products are proportional to the log co-occurrence probability. This score method assesses that directly using the current `self.embedding`. Parameters ---------- X : pd.DataFrame or np.array, shape `(self.n_words, self.n_vocab)` The original count matrix. Returns ------- float The Pearson correlation. """ X = self.convert_input_to_array(X) G = (self.model.W + self.model.C).detach().cpu().numpy() mask = X > 0 M = G.dot(G.T) X_log = utils.log_of_array_ignoring_zeros(X) row_log_prob = np.log(X.sum(axis=1)) row_log_prob = np.outer(row_log_prob, np.ones(X.shape[1])) prob = X_log - row_log_prob return np.corrcoef(prob[mask].ravel(), M[mask].ravel())[0, 1] def convert_input_to_array(self, X): if isinstance(X, pd.DataFrame): X = X.values return X @staticmethod def convert_output(X_pred, X): if isinstance(X, pd.DataFrame): X_pred = pd.DataFrame(X_pred, index=X.index) return X_pred def simple_example(): import utils utils.fix_random_seeds() X = np.array([ [4., 4., 2., 0.], [4., 61., 8., 18.], [2., 8., 10., 0.], [0., 18., 0., 5.]]) mod = TorchGloVe(embed_dim=2, max_iter=1000) print(mod) G = mod.fit(X) print("\nLearned vectors:") print(G) print("We expect the dot product of learned vectors " "to be proportional to the log co-occurrence probs. " "Let's see how close we came:") corr = mod.score(X) print("Pearson's R: {} ".format(corr)) return corr if __name__ == '__main__': simple_example()