import numpy as np from operator import itemgetter import torch import torch.nn as nn import torch.utils.data from torch_model_base import TorchModelBase import utils __author__ = "Christopher Potts" __version__ = "CS224u, Stanford, Spring 2022" class TorchRNNDataset(torch.utils.data.Dataset): def __init__(self, sequences, seq_lengths, y=None): """ Dataset class for RNN classifiers. The heavy-lifting is done by `collate_fn`, which handles the padding and packing necessary to efficiently process variable length sequences. Parameters ---------- sequences : list of `torch.LongTensor`, `len(n_examples)` seq_lengths : torch.LongTensor, shape `(n_examples, )` y : None or torch.LongTensor, shape `(n_examples, )` If None, then we are in prediction mode. Otherwise, these are indices into the list of classes. """ assert len(sequences) == len(seq_lengths) self.sequences = sequences self.seq_lengths = seq_lengths if y is not None: assert len(sequences) == len(y) self.y = y @staticmethod def collate_fn(batch): """ Format a batch of examples for use in both training and prediction. Parameters ---------- batch : tuple of length 2 (prediction) or 3 (training) The first element is the list of input sequences. The second is the list of lengths for those sequences. The third, where present, is the list of labels. Returns ------- X : torch.Tensor, shape `(batch_size, max_batch_length)` As padded by `torch.nn.utils.rnn.pad_sequence. seq_lengths : torch.LongTensor, shape `(batch_size, )` y : torch.LongTensor, shape `(batch_size, )` Only for training. In the case where `y` cannot be turned into a Tensor, we assume it is because it is a list of variable length sequences and to use `torch.nn.utils.rnn.pad_sequence`. The hope is that this will accomodate sequence prediction. """ batch_elements = list(zip(*batch)) X = batch_elements[0] seq_lengths = batch_elements[1] X = torch.nn.utils.rnn.pad_sequence(X, batch_first=True) seq_lengths = torch.tensor(seq_lengths) if len(batch_elements) == 3: y = batch_elements[2] # We can try to accommodate the case where `y` is a sequence # loss with potentially different lengths by resorting to # padding if creating a tensor is not possible: try: y = torch.tensor(y) except (ValueError, TypeError): y = torch.nn.utils.rnn.pad_sequence(y, batch_first=True) return X, seq_lengths, y else: return X, seq_lengths def __len__(self): return len(self.sequences) def __getitem__(self, idx): if self.y is not None: return self.sequences[idx], self.seq_lengths[idx], self.y[idx] else: return self.sequences[idx], self.seq_lengths[idx] class TorchRNNModel(nn.Module): def __init__(self, vocab_size, embed_dim=50, embedding=None, use_embedding=True, rnn_cell_class=nn.LSTM, hidden_dim=50, bidirectional=False, freeze_embedding=False): """ Defines the core RNN computation graph. For an explanation of the parameters, see `TorchRNNClassifierModel`. This class handles just the RNN components of the overall classifier model. `TorchRNNClassifierModel` uses the output states to create a classifier. """ super().__init__() self.vocab_size = vocab_size self.use_embedding = use_embedding self.embed_dim = embed_dim self.hidden_dim = hidden_dim self.bidirectional = bidirectional self.freeze_embedding = freeze_embedding # Graph if self.use_embedding: self.embedding = self._define_embedding( embedding, vocab_size, self.embed_dim, self.freeze_embedding) self.embed_dim = self.embedding.embedding_dim self.rnn = rnn_cell_class( input_size=self.embed_dim, hidden_size=hidden_dim, batch_first=True, bidirectional=bidirectional) def forward(self, X, seq_lengths): if self.use_embedding: X = self.embedding(X) embs = torch.nn.utils.rnn.pack_padded_sequence( X, batch_first=True, lengths=seq_lengths.cpu(), enforce_sorted=False) outputs, state = self.rnn(embs) return outputs, state @staticmethod def _define_embedding(embedding, vocab_size, embed_dim, freeze_embedding): if embedding is None: emb = nn.Embedding(vocab_size, embed_dim) emb.weight.requires_grad = not freeze_embedding return emb elif isinstance(embedding, np.ndarray): embedding = torch.FloatTensor(embedding) return nn.Embedding.from_pretrained( embedding, freeze=freeze_embedding) else: return embedding class TorchRNNClassifierModel(nn.Module): def __init__(self, rnn, output_dim, classifier_activation): """ Defines the core computation graph for `TorchRNNClassifier`. This involves using the outputs of a `TorchRNNModel` instance to build a softmax classifier: h[t] = rnn(x[t], h[t-1]) h = f(h[-1].dot(W_hy) + b_h) y = softmax(hW + b_y) This class uses its `rnn` parameter to compute each `h[1]`, and then it adds the classifier parameters that use `h[-1]` as inputs. Where `bidirectional=True`, `h[-1]` is `torch.cat([h[0], h[-1])`. """ super().__init__() self.rnn = rnn self.output_dim = output_dim self.hidden_dim = self.rnn.hidden_dim if self.rnn.bidirectional: self.classifier_dim = self.hidden_dim * 2 else: self.classifier_dim = self.hidden_dim self.hidden_layer = nn.Linear( self.classifier_dim, self.hidden_dim) self.classifier_activation = classifier_activation self.classifier_layer = nn.Linear( self.hidden_dim, self.output_dim) def forward(self, X, seq_lengths): outputs, state = self.rnn(X, seq_lengths) state = self.get_batch_final_states(state) if self.rnn.bidirectional: state = torch.cat((state[0], state[1]), dim=1) h = self.classifier_activation(self.hidden_layer(state)) logits = self.classifier_layer(h) return logits def get_batch_final_states(self, state): if self.rnn.rnn.__class__.__name__ == 'LSTM': return state[0].squeeze(0) else: return state.squeeze(0) class TorchRNNClassifier(TorchModelBase): def __init__(self, vocab, hidden_dim=50, embedding=None, use_embedding=True, embed_dim=50, rnn_cell_class=nn.LSTM, bidirectional=False, freeze_embedding=False, classifier_activation=nn.ReLU(), **base_kwargs): """ RNN-based Recurrent Neural Network for classification problems. The network will work for any kind of classification task. Parameters ---------- vocab : list of str This should be the vocabulary. It needs to be aligned with `embedding` in the sense that the ith element of vocab should be represented by the ith row of `embedding`. Ignored if `use_embedding=False`. embedding : np.array or None Each row represents a word in `vocab`, as described above. use_embedding : bool If True, then incoming examples are presumed to be lists of elements of the vocabulary. If False, then they are presumed to be lists of vectors. In this case, the `embedding` and `embed_dim` arguments are ignored, since no embedding is needed and `embed_dim` is set by the nature of the incoming vectors. embed_dim : int Dimensionality for the initial embeddings. This is ignored if `embedding` is not None, as a specified value there determines this value. Also ignored if `use_embedding=False`. rnn_cell_class : class for PyTorch recurrent layer Should be just the class name, not an instance of the class. hidden_dim : int Dimensionality of the hidden layer in the RNN. bidirectional : bool If True, then the final hidden states from passes in both directions are used. freeze_embedding : bool If True, the embedding will be updated during training. If False, the embedding will be frozen. This parameter applies to both randomly initialized and pretrained embeddings. classifier_activation : nn.Module The non-activation function used by the network for the hidden layer of the classifier. **base_kwargs For details, see `torch_model_base.py`. Attributes ---------- loss: nn.CrossEntropyLoss(reduction="mean") self.params: list Extends TorchModelBase.params with names for all of the arguments for this class to support tuning of these values using `sklearn.model_selection` tools. """ self.vocab = vocab self.hidden_dim = hidden_dim self.embedding = embedding self.use_embedding = use_embedding self.embed_dim = embed_dim self.rnn_cell_class = rnn_cell_class self.bidirectional = bidirectional self.freeze_embedding = freeze_embedding self.classifier_activation = classifier_activation super().__init__(**base_kwargs) self.params += [ 'hidden_dim', 'embed_dim', 'embedding', 'use_embedding', 'rnn_cell_class', 'bidirectional', 'freeze_embedding', 'classifier_activation'] self.loss = nn.CrossEntropyLoss(reduction="mean") def build_graph(self): """ The core computation graph. This is called by `fit`, which sets the `self.model` attribute. Returns ------- TorchRNNModel """ rnn = TorchRNNModel( vocab_size=len(self.vocab), embedding=self.embedding, use_embedding=self.use_embedding, embed_dim=self.embed_dim, rnn_cell_class=self.rnn_cell_class, hidden_dim=self.hidden_dim, bidirectional=self.bidirectional, freeze_embedding=self.freeze_embedding) model = TorchRNNClassifierModel( rnn=rnn, output_dim=self.n_classes_, classifier_activation=self.classifier_activation) self.embed_dim = rnn.embed_dim return model def build_dataset(self, X, y=None): """ Format data for training and prediction. Parameters ---------- X : list of lists The raw sequences. The lists are expected to contain elements of `self.vocab`. This method converts them to indices for PyTorch. y : list or None The raw labels. This method turns them into indices for PyTorch processing. If None, then we are in prediction mode. Returns ------- TorchRNNDataset """ X, seq_lengths = self._prepare_sequences(X) if y is None: return TorchRNNDataset(X, seq_lengths) else: self.classes_ = sorted(set(y)) self.n_classes_ = len(self.classes_) class2index = dict(zip(self.classes_, range(self.n_classes_))) y = [class2index[label] for label in y] return TorchRNNDataset(X, seq_lengths, y) def _prepare_sequences(self, X): """ Internal method for turning X into a list of indices into `self.vocab` and calculating the true lengths of the elements in `X`. Parameters ---------- X : list of lists, `len(n_examples)` Returns ------- new_X : list of lists, `len(n_examples)` seq_lengths : torch.LongTensor, shape `(n_examples, )` """ if self.use_embedding: new_X = [] seq_lengths = [] index = dict(zip(self.vocab, range(len(self.vocab)))) unk_index = index['$UNK'] for ex in X: seq = [index.get(w, unk_index) for w in ex] seq = torch.tensor(seq) new_X.append(seq) seq_lengths.append(len(seq)) else: new_X = [torch.FloatTensor(ex) for ex in X] seq_lengths = [len(ex) for ex in X] self.embed_dim = X[0][0].shape[0] seq_lengths = torch.tensor(seq_lengths) return new_X, seq_lengths def score(self, X, y, device=None): """ Uses macro-F1 as the score function. Note: this departs from `sklearn`, where classifiers use accuracy as their scoring function. Using macro-F1 is more consistent with our course. This function can be used to evaluate models, but its primary use is in cross-validation and hyperparameter tuning. Parameters ---------- X: np.array, shape `(n_examples, n_features)` y: iterable, shape `len(n_examples)` These can be the raw labels. They will converted internally as needed. See `build_dataset`. device: str or None Allows the user to temporarily change the device used during prediction. This is useful if predictions require a lot of memory and so are better done on the CPU. After prediction is done, the model is returned to `self.device`. Returns ------- float """ preds = self.predict(X, device=device) return utils.safe_macro_f1(y, preds) def predict_proba(self, X, device=None): """ Predicted probabilities for the examples in `X`. Parameters ---------- X : np.array, shape `(n_examples, n_features)` device: str or None Allows the user to temporarily change the device used during prediction. This is useful if predictions require a lot of memory and so are better done on the CPU. After prediction is done, the model is returned to `self.device`. Returns ------- np.array, shape `(len(X), self.n_classes_)` Each row of this matrix will sum to 1.0. """ preds = self._predict(X, device=device) probs = torch.softmax(preds, dim=1).cpu().numpy() return probs def predict(self, X, device=None): """ Predicted labels for the examples in `X`. These are converted from the integers that PyTorch needs back to their original values in `self.classes_`. Parameters ---------- X : np.array, shape `(n_examples, n_features)` device: str or None Allows the user to temporarily change the device used during prediction. This is useful if predictions require a lot of memory and so are better done on the CPU. After prediction is done, the model is returned to `self.device`. Returns ------- list, length len(X) """ probs = self.predict_proba(X, device=device) return [self.classes_[i] for i in probs.argmax(axis=1)] def simple_example(): utils.fix_random_seeds() vocab = ['a', 'b', '$UNK'] # No b before an a train = [ [list('ab'), 'good'], [list('aab'), 'good'], [list('abb'), 'good'], [list('aabb'), 'good'], [list('ba'), 'bad'], [list('baa'), 'bad'], [list('bba'), 'bad'], [list('bbaa'), 'bad'], [list('aba'), 'bad']] test = [ [list('baaa'), 'bad'], [list('abaa'), 'bad'], [list('bbaa'), 'bad'], [list('aaab'), 'good'], [list('aaabb'), 'good']] X_train, y_train = zip(*train) X_test, y_test = zip(*test) mod = TorchRNNClassifier(vocab) print(mod) mod.fit(X_train, y_train) preds = mod.predict(X_test) print("\nPredictions:") for ex, pred, gold in zip(X_test, preds, y_test): score = "correct" if pred == gold else "incorrect" print("{0:>6} - predicted: {1:>4}; actual: {2:>4} - {3}".format( "".join(ex), pred, gold, score)) return mod.score(X_test, y_test) if __name__ == '__main__': simple_example()