Turn BidirectionalLM into a more-general LanguageModel class (#2264)

Fixes #2255

This PR replaces the `BidirectionalLM` class with a more-general `LanguageModel` that can be used in either the unidirectional/forward setting or the bidirectional setting.

It also accordingly replaces the `BidirectionalLanguageModelTokenEmbedder` with a `LanguageModelTokenEmbedder`.

Also fixes bug in the experiment_unsampled.jsonnet config that was preventing a test from actually being unsampled.

TODO:

- [x] test the unidirectional case
- [x] properly deprecate `BidirectionalLM` and `BidirectionalLanguageModelTokenEmbedder` 
- [x] check docs for accuracy
- [x] fix user-facing training configs

Author: nelson-liu

allennlp/models/__init__.py

@@ -27,3 +27,4 @@
 from allennlp.models.bimpm import BiMpm
 from allennlp.models.graph_parser import GraphParser
 from allennlp.models.bidirectional_lm import BidirectionalLanguageModel
+from allennlp.models.language_model import LanguageModel

allennlp/models/bidirectional_lm.py

@@ -1,53 +1,16 @@
-from typing import Dict, List, Tuple, Union
+from typing import Union
 
-import torch
-import numpy as np
-
-from allennlp.common.checks import ConfigurationError
 from allennlp.data.vocabulary import Vocabulary
+from allennlp.models.language_model import LanguageModel
 from allennlp.models.model import Model
 from allennlp.modules.text_field_embedders import TextFieldEmbedder
-from allennlp.modules.sampled_softmax_loss import SampledSoftmaxLoss
 from allennlp.modules.seq2seq_encoders import Seq2SeqEncoder
-from allennlp.nn.util import get_text_field_mask
 from allennlp.nn import InitializerApplicator
 
 
-class _SoftmaxLoss(torch.nn.Module):
-    """
-    Given some embeddings and some targets, applies a linear layer
-    to create logits over possible words and then returns the
-    negative log likelihood.
-    """
-    def __init__(self,
-                 num_words: int,
-                 embedding_dim: int) -> None:
-        super().__init__()
-
-        # TODO(joelgrus): implement tie_embeddings (maybe)
-        self.tie_embeddings = False
-
-        self.softmax_w = torch.nn.Parameter(
-                torch.randn(embedding_dim, num_words) / np.sqrt(embedding_dim)
-        )
-        self.softmax_b = torch.nn.Parameter(torch.zeros(num_words))
-
-    def forward(self, embeddings: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
-        # pylint: disable=arguments-differ
-        # embeddings is size (n, embedding_dim)
-        # targets is (batch_size, ) with the correct class id
-        # Does not do any count normalization / divide by batch size
-        probs = torch.nn.functional.log_softmax(
-                torch.matmul(embeddings, self.softmax_w) + self.softmax_b,
-                dim=-1
-        )
-
-        return torch.nn.functional.nll_loss(probs, targets.long(), reduction="sum")
-
-
 @Model.register('bidirectional-language-model')
 @Model.register('bidirectional_language_model')
-class BidirectionalLanguageModel(Model):
+class BidirectionalLanguageModel(LanguageModel):
     """
     The ``BidirectionalLanguageModel`` applies a bidirectional "contextualizing"
     ``Seq2SeqEncoder`` to uncontextualized embeddings, using a ``SoftmaxLoss``
@@ -90,211 +53,12 @@ def __init__(self,
                  num_samples: int = None,
                  sparse_embeddings: bool = False,
                  initializer: InitializerApplicator = None) -> None:
-        super().__init__(vocab)
-        self._text_field_embedder = text_field_embedder
-
-        if not contextualizer.is_bidirectional():
-            raise ConfigurationError("contextualizer must be bidirectional")
-
-        self._contextualizer = contextualizer
-        # The dimension for making predictions just in the forward
-        # (or backward) direction.
-        self._forward_dim = contextualizer.get_output_dim() // 2
-
-        # TODO(joelgrus): more sampled softmax configuration options, as needed.
-        if num_samples is not None:
-            self._softmax_loss = SampledSoftmaxLoss(num_words=vocab.get_vocab_size(),
-                                                    embedding_dim=self._forward_dim,
-                                                    num_samples=num_samples,
-                                                    sparse=sparse_embeddings)
-        else:
-            self._softmax_loss = _SoftmaxLoss(num_words=vocab.get_vocab_size(),
-                                              embedding_dim=self._forward_dim)
-
-        # TODO(brendanr): Output perplexity here. e^loss
-        self.register_buffer('_last_average_loss', torch.zeros(1))
-
-        if dropout:
-            self._dropout = torch.nn.Dropout(dropout)
-        else:
-            self._dropout = lambda x: x
-
-        self._loss_scale = loss_scale
-        if initializer is not None:
-            initializer(self)
-
-    def _get_target_token_embedding(self,
-                                    token_embeddings: torch.Tensor,
-                                    mask: torch.Tensor,
-                                    direction: int) -> torch.Tensor:
-        # Need to shift the mask in the correct direction
-        zero_col = token_embeddings.new_zeros(mask.size(0), 1).byte()
-        if direction == 0:
-            # forward direction, get token to right
-            shifted_mask = torch.cat([zero_col, mask[:, 0:-1]], dim=1)
-        else:
-            shifted_mask = torch.cat([mask[:, 1:], zero_col], dim=1)
-        return token_embeddings.masked_select(shifted_mask.unsqueeze(-1)).view(-1, self._forward_dim)
-
-    def _compute_loss(self,
-                      lm_embeddings: torch.Tensor,
-                      token_embeddings: torch.Tensor,
-                      forward_targets: torch.Tensor,
-                      backward_targets: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
-        # lm_embeddings is shape (batch_size, timesteps, dim * 2)
-        # forward_targets, backward_targets are shape (batch_size, timesteps)
-        # masked with 0
-        forward_embeddings, backward_embeddings = lm_embeddings.chunk(2, -1)
-        losses: List[torch.Tensor] = []
-        for idx, embedding, targets in ((0, forward_embeddings, forward_targets),
-                                        (1, backward_embeddings, backward_targets)):
-            mask = targets > 0
-            # we need to subtract 1 to undo the padding id since the softmax
-            # does not include a padding dimension
-
-            # shape (batch_size * timesteps, )
-            non_masked_targets = targets.masked_select(mask) - 1
-
-            # shape (batch_size * timesteps, embedding_dim)
-            non_masked_embedding = embedding.masked_select(
-                    mask.unsqueeze(-1)
-            ).view(-1, self._forward_dim)
-            # note: need to return average loss across forward and backward
-            # directions, but total sum loss across all batches.
-            # Assuming batches include full sentences, forward and backward
-            # directions have the same number of samples, so sum up loss
-            # here then divide by 2 just below
-            if not self._softmax_loss.tie_embeddings or not self._use_character_inputs:
-                losses.append(self._softmax_loss(non_masked_embedding, non_masked_targets))
-            else:
-                # we also need the token embeddings corresponding to the
-                # the targets
-                raise NotImplementedError("This requires SampledSoftmaxLoss, which isn't implemented yet.")
-                # pylint: disable=unreachable
-                non_masked_token_embedding = self._get_target_token_embedding(token_embeddings, mask, idx)
-                losses.append(self._softmax(non_masked_embedding,
-                                            non_masked_targets,
-                                            non_masked_token_embedding))
-
-        return losses[0], losses[1]
-
-    def delete_softmax(self) -> None:
-        """
-        Remove the softmax weights. Useful for saving memory when calculating the loss
-        is not necessary, e.g. in an embedder.
-        """
-        self._softmax_loss = None
-
-    def num_layers(self) -> int:
-        """
-        Returns the depth of this LM. That is, how many layers the contextualizer has plus one for
-        the non-contextual layer.
-        """
-        if hasattr(self._contextualizer, 'num_layers'):
-            return self._contextualizer.num_layers + 1
-        else:
-            raise NotImplementedError(f"Contextualizer of type {type(self._contextualizer)} " +
-                                      "does not report how many layers it has.")
-
-    def forward(self,  # type: ignore
-                source: Dict[str, torch.LongTensor]) -> Dict[str, torch.Tensor]:
-        """
-        Computes the averaged forward and backward LM loss from the batch.
-
-        By convention, the input dict is required to have at least a ``"tokens"``
-        entry that's the output of a ``SingleIdTokenIndexer``, which is used
-        to compute the language model targets.
-
-        Parameters
-        ----------
-        tokens: ``torch.Tensor``, required.
-            The output of ``Batch.as_tensor_dict()`` for a batch of sentences.
-
-        Returns
-        -------
-        Dict with keys:
-
-        ``'loss'``: ``torch.Tensor``
-            averaged forward/backward negative log likelihood
-        ``'forward_loss'``: ``torch.Tensor``
-            forward direction negative log likelihood
-        ``'backward_loss'``: ``torch.Tensor``
-            backward direction negative log likelihood
-        ``'lm_embeddings'``: ``Union[torch.Tensor, List[torch.Tensor]]``
-            (batch_size, timesteps, embed_dim) tensor of top layer contextual representations or
-            list of all layers. No dropout applied.
-        ``'noncontextual_token_embeddings'``: ``torch.Tensor``
-            (batch_size, timesteps, token_embed_dim) tensor of bottom layer noncontextual
-            representations
-        ``'mask'``: ``torch.Tensor``
-            (batch_size, timesteps) mask for the embeddings
-        """
-        # pylint: disable=arguments-differ
-        mask = get_text_field_mask(source)
-
-        # shape (batch_size, timesteps, embedding_size)
-        embeddings = self._text_field_embedder(source)
-
-        # Either the top layer or all layers.
-        contextual_embeddings: Union[torch.Tensor, List[torch.Tensor]] = self._contextualizer(
-                embeddings, mask
-        )
-
-        return_dict = {}
-
-        # If we have target tokens, calculate the loss.
-        token_ids = source.get("tokens")
-        if token_ids is not None:
-            assert isinstance(contextual_embeddings, torch.Tensor)
-
-            # Use token_ids to compute targets
-            forward_targets = torch.zeros_like(token_ids)
-            backward_targets = torch.zeros_like(token_ids)
-            forward_targets[:, 0:-1] = token_ids[:, 1:]
-            backward_targets[:, 1:] = token_ids[:, 0:-1]
-
-            # add dropout
-            contextual_embeddings_with_dropout = self._dropout(contextual_embeddings)
-
-            # compute softmax loss
-            forward_loss, backward_loss = self._compute_loss(contextual_embeddings_with_dropout,
-                                                             embeddings,
-                                                             forward_targets,
-                                                             backward_targets)
-
-            num_targets = torch.sum((forward_targets > 0).long())
-            if num_targets > 0:
-                average_loss = 0.5 * (forward_loss + backward_loss) / num_targets.float()
-            else:
-                average_loss = torch.tensor(0.0).to(forward_targets.device)  # pylint: disable=not-callable
-            # this is stored to compute perplexity if needed
-            self._last_average_loss[0] = average_loss.detach().item()
-
-            if num_targets > 0:
-                # loss is directly minimized
-                if self._loss_scale == 'n_samples':
-                    scale_factor = num_targets.float()
-                else:
-                    scale_factor = self._loss_scale
-
-                return_dict.update({
-                        'loss': average_loss * scale_factor,
-                        'forward_loss': forward_loss * scale_factor / num_targets.float(),
-                        'backward_loss': backward_loss * scale_factor / num_targets.float()
-                })
-            else:
-                # average_loss zero tensor, return it for all
-                return_dict.update({
-                        'loss': average_loss,
-                        'forward_loss': average_loss,
-                        'backward_loss': average_loss
-                })
-
-        return_dict.update({
-                # Note: These embeddings do not have dropout applied.
-                'lm_embeddings': contextual_embeddings,
-                'noncontextual_token_embeddings': embeddings,
-                'mask': mask
-        })
-
-        return return_dict
+        super().__init__(vocab=vocab,
+                         text_field_embedder=text_field_embedder,
+                         contextualizer=contextualizer,
+                         dropout=dropout,
+                         loss_scale=loss_scale,
+                         num_samples=num_samples,
+                         sparse_embeddings=sparse_embeddings,
+                         bidirectional=True,
+                         initializer=initializer)