enabled lmlstm labels exporting

Summary: Enables label names to appear in the exported LM caffe2 model, which removes dependencies on external vocab files.

Differential Revision: D16190230

fbshipit-source-id: f42b8224a7eac13b4ef899a4419dc03de31bcf20

pytext/models/output_layers/lm_output_layer.py

@@ -5,12 +5,14 @@
 
 import torch
 import torch.nn.functional as F
+from caffe2.python import core
 from pytext.config.component import create_loss
 from pytext.data.utils import PAD, Vocabulary
 from pytext.fields import FieldMeta
 from pytext.loss import CrossEntropyLoss, Loss
 
 from .output_layer_base import OutputLayerBase
+from .utils import OutputLayerUtils
 
 
 class LMOutputLayer(OutputLayerBase):
@@ -107,6 +109,19 @@ def get_pred(
         scores = F.log_softmax(logit, 2)
         return preds, scores
 
+    def export_to_caffe2(
+        self,
+        workspace: core.workspace,
+        init_net: core.Net,
+        predict_net: core.Net,
+        model_out: torch.Tensor,
+        output_name: str,
+    ) -> List[core.BlobReference]:
+        prob_out = predict_net.Softmax(output_name, axis=model_out.dim() - 1)
+        return OutputLayerUtils.gen_additional_blobs(
+            predict_net, prob_out, model_out, output_name, self.target_names
+        )
+
     @staticmethod
     def calculate_perplexity(sequence_loss: torch.Tensor) -> torch.Tensor:
         return torch.exp(sequence_loss)