class CTM: """Class to train the contextualized topic model. This is the more general class that we are keeping to avoid braking code, user should use the two subclasses ZeroShotTM and CombinedTm to do topic modeling. :param input_size: int, dimension of input :param bert_input_size: int, dimension of input that comes from BERT embeddings :param inference_type: string, you can choose between the contextual model and the combined model :param num_topics: int, number of topic components, (default 10) :param model_type: string, 'prodLDA' or 'LDA' (default 'prodLDA') :param hidden_sizes: tuple, length = n_layers, (default (100, 100)) :param activation: string, 'softplus', 'relu', (default 'softplus') :param dropout: float, dropout to use (default 0.2) :param learn_priors: bool, make priors a learnable parameter (default True) :param batch_size: int, size of batch to use for training (default 64) :param lr: float, learning rate to use for training (default 2e-3) :param momentum: float, momentum to use for training (default 0.99) :param solver: string, optimizer 'adam' or 'sgd' (default 'adam') :param num_epochs: int, number of epochs to train for, (default 100) :param reduce_on_plateau: bool, reduce learning rate by 10x on plateau of 10 epochs (default False) :param num_data_loader_workers: int, number of data loader workers (default cpu_count). set it to 0 if you are using Windows """ def __init__(self, input_size, bert_input_size, inference_type="combined", num_topics=10, model_type='prodLDA', hidden_sizes=(100, 100), activation='softplus', dropout=0.2, learn_priors=True, batch_size=5, lr=2e-3, momentum=0.99, solver='adam', num_epochs=100, reduce_on_plateau=False, num_data_loader_workers=2): #mp.cpu_count() warnings.simplefilter('always', DeprecationWarning) if self.__class__.__name__ == "CTM": warnings.warn( "Direct call to CTM is deprecated and will be removed in version 2, use CombinedTM or ZeroShotTM", DeprecationWarning) assert isinstance(input_size, int) and input_size > 0, \ "input_size must by type int > 0." assert isinstance(num_topics, int) and input_size > 0, \ "num_topics must by type int > 0." assert model_type in ['LDA', 'prodLDA'], \ "model must be 'LDA' or 'prodLDA'." assert isinstance(hidden_sizes, tuple), \ "hidden_sizes must be type tuple." assert activation in ['softplus', 'relu'], \ "activation must be 'softplus' or 'relu'." assert dropout >= 0, "dropout must be >= 0." assert isinstance(learn_priors, bool), "learn_priors must be boolean." assert isinstance(batch_size, int) and batch_size > 0, \ "batch_size must be int > 0." assert lr > 0, "lr must be > 0." assert isinstance(momentum, float) and 0 < momentum <= 1, \ "momentum must be 0 < float <= 1." assert solver in ['adam', 'sgd'], "solver must be 'adam' or 'sgd'." assert isinstance(reduce_on_plateau, bool), \ "reduce_on_plateau must be type bool." assert isinstance(num_data_loader_workers, int) and num_data_loader_workers >= 0, \ "num_data_loader_workers must by type int >= 0. set 0 if you are using windows" self.input_size = input_size self.num_topics = num_topics self.model_type = model_type self.hidden_sizes = hidden_sizes self.activation = activation self.dropout = dropout self.learn_priors = learn_priors self.batch_size = batch_size self.lr = lr self.bert_size = bert_input_size self.momentum = momentum self.solver = solver self.num_epochs = num_epochs self.reduce_on_plateau = reduce_on_plateau self.num_data_loader_workers = num_data_loader_workers self.model = DecoderNetwork(input_size, self.bert_size, inference_type, num_topics, model_type, hidden_sizes, activation, dropout, learn_priors) self.early_stopping = None # init optimizer if self.solver == 'adam': self.optimizer = optim.Adam(self.model.parameters(), lr=lr, betas=(self.momentum, 0.99)) elif self.solver == 'sgd': self.optimizer = optim.SGD(self.model.parameters(), lr=lr, momentum=self.momentum) # init lr scheduler if self.reduce_on_plateau: self.scheduler = ReduceLROnPlateau(self.optimizer, patience=10) # performance attributes self.best_loss_train = float('inf') # training attributes self.model_dir = None self.train_data = None self.nn_epoch = None # validation attributes self.validation_data = None # learned topics self.best_components = None # Use cuda if available if torch.cuda.is_available(): self.USE_CUDA = True else: self.USE_CUDA = False if self.USE_CUDA: self.model = self.model.cuda() def _loss(self, inputs, word_dists, prior_mean, prior_variance, posterior_mean, posterior_variance, posterior_log_variance): # KL term # var division term var_division = torch.sum(posterior_variance / prior_variance, dim=1) # diff means term diff_means = prior_mean - posterior_mean diff_term = torch.sum((diff_means * diff_means) / prior_variance, dim=1) # logvar det division term logvar_det_division = \ prior_variance.log().sum() - posterior_log_variance.sum(dim=1) # combine terms KL = 0.5 * (var_division + diff_term - self.num_topics + logvar_det_division) # Reconstruction term RL = -torch.sum(inputs * torch.log(word_dists + 1e-10), dim=1) loss = KL + RL return loss.sum() def _train_epoch(self, loader): """Train epoch.""" self.model.train() train_loss = 0 samples_processed = 0 for batch_samples in loader: # batch_size x vocab_size X = batch_samples['X'] X = X.reshape(X.shape[0], -1) X_bert = batch_samples['X_bert'] if self.USE_CUDA: X = X.cuda() X_bert = X_bert.cuda() # forward pass self.model.zero_grad() prior_mean, prior_variance, posterior_mean, posterior_variance, posterior_log_variance, word_dists =\ self.model(X, X_bert) # backward pass loss = self._loss(X, word_dists, prior_mean, prior_variance, posterior_mean, posterior_variance, posterior_log_variance) loss.backward() self.optimizer.step() # compute train loss samples_processed += X.size()[0] train_loss += loss.item() train_loss /= samples_processed return samples_processed, train_loss def fit(self, train_dataset, validation_dataset=None, save_dir=None, verbose=False, patience=5, delta=0): """ Train the CTM model. :param train_dataset: PyTorch Dataset class for training data. :param validation_dataset: PyTorch Dataset class for validation data. If not None, the training stops if validation loss doesn't improve after a given patience :param save_dir: directory to save checkpoint models to. :param verbose: verbose :param patience: How long to wait after last time validation loss improved. Default: 5 :param delta: Minimum change in the monitored quantity to qualify as an improvement. Default: 0 """ # Print settings to output file if verbose: print("Settings: \n\ N Components: {}\n\ Topic Prior Mean: {}\n\ Topic Prior Variance: {}\n\ Model Type: {}\n\ Hidden Sizes: {}\n\ Activation: {}\n\ Dropout: {}\n\ Learn Priors: {}\n\ Learning Rate: {}\n\ Momentum: {}\n\ Reduce On Plateau: {}\n\ Save Dir: {}".format( self.num_topics, 0.0, 1. - (1. / self.num_topics), self.model_type, self.hidden_sizes, self.activation, self.dropout, self.learn_priors, self.lr, self.momentum, self.reduce_on_plateau, save_dir)) self.model_dir = save_dir self.train_data = train_dataset self.validation_data = validation_dataset if self.validation_data is not None: self.early_stopping = EarlyStopping(patience=patience, verbose=verbose, path=save_dir, delta=delta) train_loader = DataLoader(self.train_data, batch_size=self.batch_size, shuffle=True, num_workers=self.num_data_loader_workers) # init training variables train_loss = 0 samples_processed = 0 # train loop pbar = tqdm(self.num_epochs, position=0, leave=True) for epoch in range(self.num_epochs): self.nn_epoch = epoch # train epoch s = datetime.datetime.now() sp, train_loss = self._train_epoch(train_loader) samples_processed += sp e = datetime.datetime.now() pbar.update(1) pbar.set_description( "Epoch: [{}/{}]\t Seen Samples: [{}/{}]\tTrain Loss: {}\tTime: {}" .format(epoch + 1, self.num_epochs, samples_processed, len(self.train_data) * self.num_epochs, train_loss, e - s)) if self.validation_data is not None: self.best_loss_train = train_loss self.best_components = self.model.beta validation_loader = DataLoader( self.validation_data, batch_size=self.batch_size, shuffle=True, num_workers=self.num_data_loader_workers) # train epoch s = datetime.datetime.now() val_samples_processed, val_loss = self._validation( validation_loader) e = datetime.datetime.now() # report if verbose: print( "Epoch: [{}/{}]\tSamples: [{}/{}]\tValidation Loss: {}\tTime: {}" .format(epoch + 1, self.num_epochs, val_samples_processed, len(self.validation_data) * self.num_epochs, val_loss, e - s)) self.early_stopping(val_loss, self) if self.early_stopping.early_stop: print("Early stopping") break else: # save best if train_loss < self.best_loss_train: self.best_loss_train = train_loss self.best_components = self.model.beta if save_dir is not None: self.save(save_dir) pbar.close() def _validation(self, loader): """Validation epoch.""" self.model.eval() val_loss = 0 samples_processed = 0 for batch_samples in loader: # batch_size x vocab_size X = batch_samples['X'] X = X.reshape(X.shape[0], -1) X_bert = batch_samples['X_bert'] if self.USE_CUDA: X = X.cuda() X_bert = X_bert.cuda() # forward pass self.model.zero_grad() prior_mean, prior_variance, posterior_mean, posterior_variance, posterior_log_variance, word_dists =\ self.model(X, X_bert) loss = self._loss(X, word_dists, prior_mean, prior_variance, posterior_mean, posterior_variance, posterior_log_variance) # compute train loss samples_processed += X.size()[0] val_loss += loss.item() val_loss /= samples_processed return samples_processed, val_loss def get_thetas(self, dataset, n_samples=20): """ Get the document-topic distribution for a dataset of topics. Includes multiple sampling to reduce variation via the parameter n_sample. :param dataset: a PyTorch Dataset containing the documents :param n_samples: the number of sample to collect to estimate the final distribution (the more the better). """ warnings.warn( "Call to `get_thetas` is deprecated and will be removed in version 2, " "use `get_doc_topic_distribution` instead", DeprecationWarning) return self.get_doc_topic_distribution(dataset, n_samples=n_samples) def get_doc_topic_distribution(self, dataset, n_samples=20): """ Get the document-topic distribution for a dataset of topics. Includes multiple sampling to reduce variation via the parameter n_sample. :param dataset: a PyTorch Dataset containing the documents :param n_samples: the number of sample to collect to estimate the final distribution (the more the better). """ self.model.eval() loader = DataLoader(dataset, batch_size=self.batch_size, shuffle=False, num_workers=self.num_data_loader_workers) pbar = tqdm(n_samples, position=0, leave=True) final_thetas = [] for sample_index in range(n_samples): with torch.no_grad(): collect_theta = [] for batch_samples in loader: # batch_size x vocab_size X = batch_samples['X'] X = X.reshape(X.shape[0], -1) X_bert = batch_samples['X_bert'] if self.USE_CUDA: X = X.cuda() X_bert = X_bert.cuda() # forward pass self.model.zero_grad() collect_theta.extend( self.model.get_theta(X, X_bert).cpu().numpy().tolist()) pbar.update(1) pbar.set_description("Sampling: [{}/{}]".format( sample_index + 1, n_samples)) final_thetas.append(np.array(collect_theta)) pbar.close() return np.sum(final_thetas, axis=0) / n_samples def get_most_likely_topic(self, doc_topic_distribution): """ get the most likely topic for each document :param doc_topic_distribution: ndarray representing the topic distribution of each document """ return np.argmax(doc_topic_distribution, axis=0) def predict(self, dataset, k=9): """Predict input.""" self.model.eval() loader = DataLoader(dataset, batch_size=self.batch_size, shuffle=False, num_workers=self.num_data_loader_workers) preds = [] with torch.no_grad(): for batch_samples in loader: # batch_size x vocab_size X = batch_samples['X'] X = X.reshape(X.shape[0], -1) X_bert = batch_samples['X_bert'] if self.USE_CUDA: X = X.cuda() X_bert = X_bert.cuda() # forward pass self.model.zero_grad() _, _, _, _, _, word_dists = self.model(X, X_bert) _, indices = torch.sort(word_dists, dim=1) preds += [indices[:, :k]] preds = torch.cat(preds, dim=0) return preds def get_topics(self, k=10): """ Retrieve topic words. :param k: int, number of words to return per topic, default 10. """ assert k <= self.input_size, "k must be <= input size." component_dists = self.best_components topics = defaultdict(list) for i in range(self.num_topics): _, idxs = torch.topk(component_dists[i], k) component_words = [ self.train_data.idx2token[idx] for idx in idxs.cpu().numpy() ] topics[i] = component_words return topics def get_topic_lists(self, k=10): """ Retrieve the lists of topic words. :param k: (int) number of words to return per topic, default 10. """ assert k <= self.input_size, "k must be <= input size." # TODO: collapse this method with the one that just returns the topics component_dists = self.best_components topics = [] for i in range(self.num_topics): _, idxs = torch.topk(component_dists[i], k) component_words = [ self.train_data.idx2token[idx] for idx in idxs.cpu().numpy() ] print(component_words) topics.append(component_words) return topics def _format_file(self): model_dir = "contextualized_topic_model_nc_{}_tpm_{}_tpv_{}_hs_{}_ac_{}_do_{}_lr_{}_mo_{}_rp_{}". \ format(self.num_topics, 0.0, 1 - (1. / self.num_topics), self.model_type, self.hidden_sizes, self.activation, self.dropout, self.lr, self.momentum, self.reduce_on_plateau) return model_dir def save(self, models_dir=None): """ Save model. (Experimental Feature, not tested) :param models_dir: path to directory for saving NN models. """ warnings.simplefilter('always', Warning) warnings.warn( "This is an experimental feature that we has not been fully tested. Refer to the following issue:" "https://github.com/MilaNLProc/contextualized-topic-models/issues/38", Warning) if (self.model is not None) and (models_dir is not None): model_dir = self._format_file() if not os.path.isdir(os.path.join(models_dir, model_dir)): os.makedirs(os.path.join(models_dir, model_dir)) filename = "epoch_{}".format(self.nn_epoch) + '.pth' fileloc = os.path.join(models_dir, model_dir, filename) with open(fileloc, 'wb') as file: torch.save( { 'state_dict': self.model.state_dict(), 'dcue_dict': self.__dict__ }, file) def load(self, model_dir, epoch): """ Load a previously trained model. (Experimental Feature, not tested) :param model_dir: directory where models are saved. :param epoch: epoch of model to load. """ warnings.simplefilter('always', Warning) warnings.warn( "This is an experimental feature that we has not been fully tested. Refer to the following issue:" "https://github.com/MilaNLProc/contextualized-topic-models/issues/38", Warning) epoch_file = "epoch_" + str(epoch) + ".pth" model_file = os.path.join(model_dir, epoch_file) with open(model_file, 'rb') as model_dict: checkpoint = torch.load(model_dict) for (k, v) in checkpoint['dcue_dict'].items(): setattr(self, k, v) self.model.load_state_dict(checkpoint['state_dict']) def get_topic_word_matrix(self): """ Return the topic-word matrix (dimensions: number of topics x length of the vocabulary). If model_type is LDA, the matrix is normalized; otherwise the matrix is unnormalized. """ return self.model.topic_word_matrix.cpu().detach().numpy() def get_topic_word_distribution(self): """ Return the topic-word distribution (dimensions: number of topics x length of the vocabulary). """ mat = self.get_topic_word_matrix() return softmax(mat, axis=1) def get_word_distribution_by_topic_id(self, topic): """ Return the word probability distribution of a topic sorted by probability. :param topic: id of the topic (int) :returns list of tuples (word, probability) sorted by the probability in descending order """ if topic >= self.num_topics: raise Exception('Topic id must be lower than the number of topics') else: wd = self.get_topic_word_distribution() t = [(word, wd[topic][idx]) for idx, word in self.train_data.idx2token.items()] t = sorted(t, key=lambda x: -x[1]) return t def get_wordcloud(self, topic_id, n_words=5, background_color="black"): """ Plotting the wordcloud. It is an adapted version of the code found here: http://amueller.github.io/word_cloud/auto_examples/simple.html#sphx-glr-auto-examples-simple-py and here https://github.com/ddangelov/Top2Vec/blob/master/top2vec/Top2Vec.py :param topic_id: id of the topic :param n_words: number of words to show in word cloud :param background_color: color of the background """ word_score_list = self.get_word_distribution_by_topic_id( topic_id)[:n_words] word_score_dict = {tup[0]: tup[1] for tup in word_score_list} plt.figure(figsize=(10, 4), dpi=200) plt.axis("off") plt.imshow( wordcloud.WordCloud( width=900, height=350, background_color=background_color).generate_from_frequencies( word_score_dict)) plt.title("Displaying Topic " + str(topic_id), loc='center', fontsize=22) plt.show() def get_predicted_topics(self, dataset, n_samples): """ Return the a list containing the predicted topic for each document (length: number of documents). :param dataset: CTMDataset to infer topics :param n_samples: number of sampling of theta :return: the predicted topics """ predicted_topics = [] thetas = self.get_doc_topic_distribution(dataset, n_samples) for idd in range(len(dataset)): predicted_topic = np.argmax(thetas[idd] / np.sum(thetas[idd])) predicted_topics.append(predicted_topic) return predicted_topics
class CTM(object): """Class to train CTM model.""" def __init__(self, input_size, bert_input_size, inference_type, n_components=10, model_type='prodLDA', hidden_sizes=(100, 100), activation='softplus', dropout=0.2, learn_priors=True, batch_size=64, lr=2e-3, momentum=0.99, solver='adam', num_epochs=100, reduce_on_plateau=False, num_data_loader_workers=mp.cpu_count()): """ Initialize CTM model. Args input_size : int, dimension of input n_components : int, number of topic components, (default 10) model_type : string, 'prodLDA' or 'LDA' (default 'prodLDA') hidden_sizes : tuple, length = n_layers, (default (100, 100)) activation : string, 'softplus', 'relu', (default 'softplus') dropout : float, dropout to use (default 0.2) learn_priors : bool, make priors a learnable parameter (default True) batch_size : int, size of batch to use for training (default 64) lr : float, learning rate to use for training (default 2e-3) momentum : float, momentum to use for training (default 0.99) solver : string, optimizer 'adam' or 'sgd' (default 'adam') num_epochs : int, number of epochs to train for, (default 100) reduce_on_plateau : bool, reduce learning rate by 10x on plateau of 10 epochs (default False) num_data_loader_workers : int, number of data loader workers (default cpu_count). set it to 0 if you are using Windows """ assert isinstance(input_size, int) and input_size > 0,\ "input_size must by type int > 0." assert isinstance(n_components, int) and input_size > 0,\ "n_components must by type int > 0." assert model_type in ['LDA', 'prodLDA'],\ "model must be 'LDA' or 'prodLDA'." assert isinstance(hidden_sizes, tuple), \ "hidden_sizes must be type tuple." assert activation in ['softplus', 'relu'], \ "activation must be 'softplus' or 'relu'." assert dropout >= 0, "dropout must be >= 0." assert isinstance(learn_priors, bool), "learn_priors must be boolean." assert isinstance(batch_size, int) and batch_size > 0,\ "batch_size must be int > 0." assert lr > 0, "lr must be > 0." assert isinstance(momentum, float) and momentum > 0 and momentum <= 1,\ "momentum must be 0 < float <= 1." assert solver in ['adam', 'sgd'], "solver must be 'adam' or 'sgd'." assert isinstance(reduce_on_plateau, bool),\ "reduce_on_plateau must be type bool." assert isinstance(num_data_loader_workers, int) and num_data_loader_workers >= 0, \ "num_data_loader_workers must by type int >= 0. set 0 if you are using windows" self.input_size = input_size self.n_components = n_components self.model_type = model_type self.hidden_sizes = hidden_sizes self.activation = activation self.dropout = dropout self.learn_priors = learn_priors self.batch_size = batch_size self.lr = lr self.bert_size = bert_input_size self.momentum = momentum self.solver = solver self.num_epochs = num_epochs self.reduce_on_plateau = reduce_on_plateau self.num_data_loader_workers = num_data_loader_workers # init inference avitm network self.model = DecoderNetwork(input_size, self.bert_size, inference_type, n_components, model_type, hidden_sizes, activation, dropout, learn_priors) # init optimizer if self.solver == 'adam': self.optimizer = optim.Adam(self.model.parameters(), lr=lr, betas=(self.momentum, 0.99)) elif self.solver == 'sgd': self.optimizer = optim.SGD(self.model.parameters(), lr=lr, momentum=self.momentum) # init lr scheduler if self.reduce_on_plateau: self.scheduler = ReduceLROnPlateau(self.optimizer, patience=10) # performance attributes self.best_loss_train = float('inf') # training atributes self.model_dir = None self.train_data = None self.nn_epoch = None # learned topics self.best_components = None # Use cuda if available if torch.cuda.is_available(): self.USE_CUDA = True else: self.USE_CUDA = False if self.USE_CUDA: self.model = self.model.cuda() def _loss(self, inputs, word_dists, prior_mean, prior_variance, posterior_mean, posterior_variance, posterior_log_variance): # KL term # var division term var_division = torch.sum(posterior_variance / prior_variance, dim=1) # diff means term diff_means = prior_mean - posterior_mean diff_term = torch.sum((diff_means * diff_means) / prior_variance, dim=1) # logvar det division term logvar_det_division = \ prior_variance.log().sum() - posterior_log_variance.sum(dim=1) # combine terms KL = 0.5 * (var_division + diff_term - self.n_components + logvar_det_division) # Reconstruction term RL = -torch.sum(inputs * torch.log(word_dists + 1e-10), dim=1) loss = KL + RL return loss.sum() def _train_epoch(self, loader): """Train epoch.""" self.model.train() train_loss = 0 samples_processed = 0 for batch_samples in loader: # batch_size x vocab_size X = batch_samples['X'] X = X.reshape(X.shape[0], -1) X_bert = batch_samples['X_bert'] if self.USE_CUDA: X = X.cuda() X_bert = X_bert.cuda() # forward pass self.model.zero_grad() prior_mean, prior_variance, \ posterior_mean, posterior_variance, posterior_log_variance, \ word_dists = self.model(X, X_bert) # backward pass loss = self._loss(X, word_dists, prior_mean, prior_variance, posterior_mean, posterior_variance, posterior_log_variance) loss.backward() self.optimizer.step() # compute train loss samples_processed += X.size()[0] train_loss += loss.item() train_loss /= samples_processed return samples_processed, train_loss def fit(self, train_dataset, save_dir=None): """ Train the AVITM model. Args train_dataset : PyTorch Dataset classs for training data. val_dataset : PyTorch Dataset classs for validation data. save_dir : directory to save checkpoint models to. """ # Print settings to output file print("Settings: \n\ N Components: {}\n\ Topic Prior Mean: {}\n\ Topic Prior Variance: {}\n\ Model Type: {}\n\ Hidden Sizes: {}\n\ Activation: {}\n\ Dropout: {}\n\ Learn Priors: {}\n\ Learning Rate: {}\n\ Momentum: {}\n\ Reduce On Plateau: {}\n\ Save Dir: {}".format(self.n_components, 0.0, 1. - (1. / self.n_components), self.model_type, self.hidden_sizes, self.activation, self.dropout, self.learn_priors, self.lr, self.momentum, self.reduce_on_plateau, save_dir)) self.model_dir = save_dir self.train_data = train_dataset train_loader = DataLoader(self.train_data, batch_size=self.batch_size, shuffle=True, num_workers=self.num_data_loader_workers) # init training variables train_loss = 0 samples_processed = 0 # train loop for epoch in range(self.num_epochs): self.nn_epoch = epoch # train epoch s = datetime.datetime.now() sp, train_loss = self._train_epoch(train_loader) samples_processed += sp e = datetime.datetime.now() # report print("Epoch: [{}/{}]\tSamples: [{}/{}]\tTrain Loss: {}\tTime: {}". format(epoch + 1, self.num_epochs, samples_processed, len(self.train_data) * self.num_epochs, train_loss, e - s)) # save best if train_loss < self.best_loss_train: self.best_loss_train = train_loss self.best_components = self.model.beta if save_dir is not None: self.save(save_dir) def get_thetas(self, dataset): """Predict input.""" self.model.eval() loader = DataLoader(dataset, batch_size=self.batch_size, shuffle=False, num_workers=self.num_data_loader_workers) with torch.no_grad(): collect_theta = [] for batch_samples in loader: # batch_size x vocab_size X = batch_samples['X'] X = X.reshape(X.shape[0], -1) X_bert = batch_samples['X_bert'] if self.USE_CUDA: X = X.cuda() X_bert = X_bert.cuda() # forward pass self.model.zero_grad() collect_theta.extend( self.model.get_theta(X, X_bert).cpu().numpy().tolist()) return collect_theta def predict(self, dataset, k=10): """Predict input.""" self.model.eval() loader = DataLoader(dataset, batch_size=self.batch_size, shuffle=False, num_workers=self.num_data_loader_workers) preds = [] with torch.no_grad(): for batch_samples in loader: # batch_size x vocab_size X = batch_samples['X'] X = X.reshape(X.shape[0], -1) X_bert = batch_samples['X_bert'] if self.USE_CUDA: X = X.cuda() X_bert = X_bert.cuda() # forward pass self.model.zero_grad() _, _, _, _, _, word_dists = self.model(X, X_bert) _, indices = torch.sort(word_dists, dim=1) preds += [indices[:, :k]] preds = torch.cat(preds, dim=0) return preds def score(self, scorer='coherence', k=10, topics=5): """Score model.""" if scorer == 'perplexity': # score = perplexity_score(truth, preds) raise NotImplementedError("Not implemented yet.") elif scorer == 'coherence': score = self._get_coherence(k, topics=topics) else: raise ValueError("Unknown score type!") return score def _get_coherence(self, k=10, topics=5): """Get coherence using palmetto web service.""" component_dists = self.best_components base_url = 'https://palmetto.demos.dice-research.org/service/cv?words=' scores = [] i = 0 while i < topics: print(i) t = np.random.randint(0, self.n_components) _, idxs = torch.topk(component_dists[t], k) component_words = [ self.train_data.idx2token[idx] for idx in idxs.cpu().numpy() ] url = base_url + '%20'.join(component_words) print(url) try: score = float(requests.get(url, timeout=300).content) scores += [score] i += 1 except requests.exceptions.Timeout: print("Attempted scoring timed out. Trying again.") continue return np.mean(scores) def get_topics(self, k=10): """ Retrieve topic words. Args k : (int) number of words to return per topic, default 10. """ assert k <= self.input_size, "k must be <= input size." component_dists = self.best_components topics = defaultdict(list) for i in range(self.n_components): _, idxs = torch.topk(component_dists[i], k) component_words = [ self.train_data.idx2token[idx] for idx in idxs.cpu().numpy() ] topics[i] = component_words return topics def get_topic_lists(self, k=10): """ Retrieve the lists of topic words. Args k : (int) number of words to return per topic, default 10. """ assert k <= self.input_size, "k must be <= input size." component_dists = self.best_components topics = [] for i in range(self.n_components): _, idxs = torch.topk(component_dists[i], k) component_words = [ self.train_data.idx2token[idx] for idx in idxs.cpu().numpy() ] topics.append(component_words) return topics def _format_file(self): model_dir = "AVITM_nc_{}_tpm_{}_tpv_{}_hs_{}_ac_{}_do_{}_lr_{}_mo_{}_rp_{}".\ format(self.n_components, 0.0, 1 - (1./self.n_components), self.model_type, self.hidden_sizes, self.activation, self.dropout, self.lr, self.momentum, self.reduce_on_plateau) return model_dir def save(self, models_dir=None): """ Save model. Args models_dir: path to directory for saving NN models. """ if (self.model is not None) and (models_dir is not None): model_dir = self._format_file() if not os.path.isdir(os.path.join(models_dir, model_dir)): os.makedirs(os.path.join(models_dir, model_dir)) filename = "epoch_{}".format(self.nn_epoch) + '.pth' fileloc = os.path.join(models_dir, model_dir, filename) with open(fileloc, 'wb') as file: torch.save( { 'state_dict': self.model.state_dict(), 'dcue_dict': self.__dict__ }, file) def load(self, model_dir, epoch): """ Load a previously trained model. Args model_dir: directory where models are saved. epoch: epoch of model to load. """ epoch_file = "epoch_" + str(epoch) + ".pth" model_file = os.path.join(model_dir, epoch_file) with open(model_file, 'rb') as model_dict: checkpoint = torch.load(model_dict) for (k, v) in checkpoint['dcue_dict'].items(): setattr(self, k, v) self._init_nn() #TODO implement this method self.model.load_state_dict(checkpoint['state_dict']) def get_topic_word_matrix(self): """ Return the topic-word matrix (dimensions: number of topics x length of the vocabulary). If model_type is LDA, the matrix is normalized; otherwise the matrix is unnormalized. """ return self.model.topic_word_matrix.cpu().detach().numpy()
class CTM(object): """Class to train the contextualized topic model. :param input_size: int, dimension of input :param bert_input_size: int, dimension of input that comes from BERT embeddings :param inference_type: string, you can choose between the contextual model and the combined model :param n_components: int, number of topic components, (default 10) :param model_type: string, 'prodLDA' or 'LDA' (default 'prodLDA') :param hidden_sizes: tuple, length = n_layers, (default (100, 100)) :param activation: string, 'softplus', 'relu', (default 'softplus') :param dropout: float, dropout to use (default 0.2) :param learn_priors: bool, make priors a learnable parameter (default True) :param batch_size: int, size of batch to use for training (default 64) :param lr: float, learning rate to use for training (default 2e-3) :param momentum: float, momentum to use for training (default 0.99) :param solver: string, optimizer 'adam' or 'sgd' (default 'adam') :param num_epochs: int, number of epochs to train for, (default 100) :param reduce_on_plateau: bool, reduce learning rate by 10x on plateau of 10 epochs (default False) :param num_data_loader_workers: int, number of data loader workers (default cpu_count). set it to 0 if you are using Windows """ def __init__(self, input_size, bert_input_size, inference_type="combined", n_components=10, model_type='prodLDA', hidden_sizes=(100, 100), activation='softplus', dropout=0.2, learn_priors=True, batch_size=64, lr=2e-3, momentum=0.99, solver='adam', num_epochs=100, reduce_on_plateau=False, num_data_loader_workers=mp.cpu_count()): warnings.simplefilter('always', DeprecationWarning) if self.__class__.__name__ == "CTM": warnings.warn( "Direct call to CTM is deprecated and will be removed in version 2, use CombinedTM or ZeroShotTM", DeprecationWarning) assert isinstance(input_size, int) and input_size > 0,\ "input_size must by type int > 0." assert isinstance(n_components, int) and input_size > 0,\ "n_components must by type int > 0." assert model_type in ['LDA', 'prodLDA'],\ "model must be 'LDA' or 'prodLDA'." assert isinstance(hidden_sizes, tuple), \ "hidden_sizes must be type tuple." assert activation in ['softplus', 'relu'], \ "activation must be 'softplus' or 'relu'." assert dropout >= 0, "dropout must be >= 0." assert isinstance(learn_priors, bool), "learn_priors must be boolean." assert isinstance(batch_size, int) and batch_size > 0,\ "batch_size must be int > 0." assert lr > 0, "lr must be > 0." assert isinstance(momentum, float) and momentum > 0 and momentum <= 1,\ "momentum must be 0 < float <= 1." assert solver in ['adam', 'sgd'], "solver must be 'adam' or 'sgd'." assert isinstance(reduce_on_plateau, bool),\ "reduce_on_plateau must be type bool." assert isinstance(num_data_loader_workers, int) and num_data_loader_workers >= 0, \ "num_data_loader_workers must by type int >= 0. set 0 if you are using windows" self.input_size = input_size self.n_components = n_components self.model_type = model_type self.hidden_sizes = hidden_sizes self.activation = activation self.dropout = dropout self.learn_priors = learn_priors self.batch_size = batch_size self.lr = lr self.bert_size = bert_input_size self.momentum = momentum self.solver = solver self.num_epochs = num_epochs self.reduce_on_plateau = reduce_on_plateau self.num_data_loader_workers = num_data_loader_workers # init inference avitm network self.model = DecoderNetwork(input_size, self.bert_size, inference_type, n_components, model_type, hidden_sizes, activation, dropout, learn_priors) # init optimizer if self.solver == 'adam': self.optimizer = optim.Adam(self.model.parameters(), lr=lr, betas=(self.momentum, 0.99)) elif self.solver == 'sgd': self.optimizer = optim.SGD(self.model.parameters(), lr=lr, momentum=self.momentum) # init lr scheduler if self.reduce_on_plateau: self.scheduler = ReduceLROnPlateau(self.optimizer, patience=10) # performance attributes self.best_loss_train = float('inf') # training atributes self.model_dir = None self.train_data = None self.nn_epoch = None # learned topics self.best_components = None # Use cuda if available if torch.cuda.is_available(): self.USE_CUDA = True else: self.USE_CUDA = False if self.USE_CUDA: self.model = self.model.cuda() def _loss(self, inputs, word_dists, prior_mean, prior_variance, posterior_mean, posterior_variance, posterior_log_variance): # KL term # var division term var_division = torch.sum(posterior_variance / prior_variance, dim=1) # diff means term diff_means = prior_mean - posterior_mean diff_term = torch.sum((diff_means * diff_means) / prior_variance, dim=1) # logvar det division term logvar_det_division = \ prior_variance.log().sum() - posterior_log_variance.sum(dim=1) # combine terms KL = 0.5 * (var_division + diff_term - self.n_components + logvar_det_division) # Reconstruction term RL = -torch.sum(inputs * torch.log(word_dists + 1e-10), dim=1) loss = KL + RL return loss.sum() def _train_epoch(self, loader): """Train epoch.""" self.model.train() train_loss = 0 samples_processed = 0 for batch_samples in loader: # batch_size x vocab_size X = batch_samples['X'] X = X.reshape(X.shape[0], -1) X_bert = batch_samples['X_bert'] if self.USE_CUDA: X = X.cuda() X_bert = X_bert.cuda() # forward pass self.model.zero_grad() prior_mean, prior_variance, \ posterior_mean, posterior_variance, posterior_log_variance, \ word_dists = self.model(X, X_bert) # backward pass loss = self._loss(X, word_dists, prior_mean, prior_variance, posterior_mean, posterior_variance, posterior_log_variance) loss.backward() self.optimizer.step() # compute train loss samples_processed += X.size()[0] train_loss += loss.item() train_loss /= samples_processed return samples_processed, train_loss def fit(self, train_dataset, save_dir=None, verbose=True): """ Train the CTM model. :param train_dataset: PyTorch Dataset class for training data. :param save_dir: directory to save checkpoint models to. """ # Print settings to output file if verbose: print("Settings: \n\ N Components: {}\n\ Topic Prior Mean: {}\n\ Topic Prior Variance: {}\n\ Model Type: {}\n\ Hidden Sizes: {}\n\ Activation: {}\n\ Dropout: {}\n\ Learn Priors: {}\n\ Learning Rate: {}\n\ Momentum: {}\n\ Reduce On Plateau: {}\n\ Save Dir: {}".format( self.n_components, 0.0, 1. - (1. / self.n_components), self.model_type, self.hidden_sizes, self.activation, self.dropout, self.learn_priors, self.lr, self.momentum, self.reduce_on_plateau, save_dir)) self.model_dir = save_dir self.train_data = train_dataset train_loader = DataLoader(self.train_data, batch_size=self.batch_size, shuffle=True, num_workers=self.num_data_loader_workers) # init training variables train_loss = 0 samples_processed = 0 # train loop for epoch in range(self.num_epochs): self.nn_epoch = epoch # train epoch s = datetime.datetime.now() sp, train_loss = self._train_epoch(train_loader) samples_processed += sp e = datetime.datetime.now() if verbose: print( "Epoch: [{}/{}]\tSamples: [{}/{}]\tTrain Loss: {}\tTime: {}" .format(epoch + 1, self.num_epochs, samples_processed, len(self.train_data) * self.num_epochs, train_loss, e - s)) # save best if train_loss < self.best_loss_train: self.best_loss_train = train_loss self.best_components = self.model.beta if save_dir is not None: self.save(save_dir) def get_thetas(self, dataset, n_samples=20): """ Get the document-topic distribution for a dataset of topics. Includes multiple sampling to reduce variation via the parameter n_sample. :param dataset: a PyTorch Dataset containing the documents :param n_samples: the number of sample to collect to estimate the final distribution (the more the better). """ self.model.eval() loader = DataLoader(dataset, batch_size=self.batch_size, shuffle=False, num_workers=self.num_data_loader_workers) final_thetas = [] for _ in range(n_samples): with torch.no_grad(): collect_theta = [] for batch_samples in loader: # batch_size x vocab_size X = batch_samples['X'] X = X.reshape(X.shape[0], -1) X_bert = batch_samples['X_bert'] if self.USE_CUDA: X = X.cuda() X_bert = X_bert.cuda() # forward pass self.model.zero_grad() collect_theta.extend( self.model.get_theta(X, X_bert).cpu().numpy().tolist()) final_thetas.append(np.array(collect_theta)) return np.sum(final_thetas, axis=0) / n_samples def predict(self, dataset, k=10): """Predict input.""" self.model.eval() loader = DataLoader(dataset, batch_size=self.batch_size, shuffle=False, num_workers=self.num_data_loader_workers) preds = [] with torch.no_grad(): for batch_samples in loader: # batch_size x vocab_size X = batch_samples['X'] X = X.reshape(X.shape[0], -1) X_bert = batch_samples['X_bert'] if self.USE_CUDA: X = X.cuda() X_bert = X_bert.cuda() # forward pass self.model.zero_grad() _, _, _, _, _, word_dists = self.model(X, X_bert) _, indices = torch.sort(word_dists, dim=1) preds += [indices[:, :k]] preds = torch.cat(preds, dim=0) return preds def get_topics(self, k=10): """ Retrieve topic words. :param k: int, number of words to return per topic, default 10. """ assert k <= self.input_size, "k must be <= input size." component_dists = self.best_components topics = defaultdict(list) for i in range(self.n_components): _, idxs = torch.topk(component_dists[i], k) component_words = [ self.train_data.idx2token[idx] for idx in idxs.cpu().numpy() ] topics[i] = component_words return topics def get_topic_lists(self, k=10): """ Retrieve the lists of topic words. :param k: (int) number of words to return per topic, default 10. """ assert k <= self.input_size, "k must be <= input size." # TODO: collapse this method with the one that just returns the topics component_dists = self.best_components topics = [] for i in range(self.n_components): _, idxs = torch.topk(component_dists[i], k) component_words = [ self.train_data.idx2token[idx] for idx in idxs.cpu().numpy() ] topics.append(component_words) return topics def _format_file(self): model_dir = "AVITM_nc_{}_tpm_{}_tpv_{}_hs_{}_ac_{}_do_{}_lr_{}_mo_{}_rp_{}".\ format(self.n_components, 0.0, 1 - (1./self.n_components), self.model_type, self.hidden_sizes, self.activation, self.dropout, self.lr, self.momentum, self.reduce_on_plateau) return model_dir def save(self, models_dir=None): """ Save model. :param models_dir: path to directory for saving NN models. """ if (self.model is not None) and (models_dir is not None): model_dir = self._format_file() if not os.path.isdir(os.path.join(models_dir, model_dir)): os.makedirs(os.path.join(models_dir, model_dir)) filename = "epoch_{}".format(self.nn_epoch) + '.pth' fileloc = os.path.join(models_dir, model_dir, filename) with open(fileloc, 'wb') as file: torch.save( { 'state_dict': self.model.state_dict(), 'dcue_dict': self.__dict__ }, file) def load(self, model_dir, epoch): """ Load a previously trained model. :param model_dir: directory where models are saved. :param epoch: epoch of model to load. """ epoch_file = "epoch_" + str(epoch) + ".pth" model_file = os.path.join(model_dir, epoch_file) with open(model_file, 'rb') as model_dict: checkpoint = torch.load(model_dict) for (k, v) in checkpoint['dcue_dict'].items(): setattr(self, k, v) self.model.load_state_dict(checkpoint['state_dict']) def get_topic_word_matrix(self): """ Return the topic-word matrix (dimensions: number of topics x length of the vocabulary). If model_type is LDA, the matrix is normalized; otherwise the matrix is unnormalized. """ return self.model.topic_word_matrix.cpu().detach().numpy() def get_predicted_topics(self, dataset, n_samples): """ Return the a list containing the predicted topic for each document (length: number of documents). :param dataset: CTMDataset to infer topics :param n_samples: number of sampling of theta :return: the predicted topics """ predicted_topics = [] #thetas = np.zeros((len(dataset), self.n_components)) thetas = self.get_thetas(dataset, n_samples) for idd in range(len(dataset)): predicted_topic = np.argmax(thetas[idd] / np.sum(thetas[idd])) predicted_topics.append(predicted_topic) return predicted_topics