use ScriptXLMTensorizer (#1123)

Summary:
Pull Request resolved: #1123

use ScriptXLMTensorizer

Reviewed By: rutyrinott

Differential Revision: D18364254

fbshipit-source-id: 1d607288f9e10c909f7a42fc2a5d53d94ca2e815

Author: chenyangyu1988

pytext/data/xlm_tensorizer.py

@@ -11,6 +11,8 @@
 from pytext.data.tokenizers import Tokenizer
 from pytext.data.utils import EOS, MASK, PAD, UNK, Vocabulary
 from pytext.data.xlm_constants import LANG2ID_15
+from pytext.torchscript.tensorizer import ScriptXLMTensorizer
+from pytext.torchscript.vocab import ScriptVocabulary
 
 
 class XLMTensorizer(BERTTensorizerBase):
@@ -85,6 +87,7 @@ def __init__(
         # unlike BERT, XLM uses the EOS token for both beginning and end of
         # sentence
         self.bos_token = self.vocab.eos_token
+        self.default_language = "en"
 
     @property
     def column_schema(self):
@@ -103,7 +106,7 @@ def get_lang_id(self, row: Dict, col: str) -> int:
             return lang_id
         else:
             # use En as default
-            return self.lang2id.get("en", 0)
+            return self.lang2id.get(self.default_language, 0)
 
     def _lookup_tokens(self, text: str, seq_len: int) -> List[str]:
         return lookup_tokens(
@@ -137,3 +140,22 @@ def numberize(self, row: Dict) -> Tuple[Any, ...]:
         seq_len = len(tokens)
         positions = [index for index in range(seq_len)]
         return tokens, segment_labels, seq_len, positions
+
+    def torchscriptify(self):
+        languages = [0] * (max(list(self.lang2id.values())) + 1)
+        for k, v in self.lang2id.items():
+            languages[v] = k
+
+        return ScriptXLMTensorizer(
+            tokenizer=self.tokenizer.torchscriptify(),
+            token_vocab=ScriptVocabulary(
+                list(self.vocab),
+                pad_idx=self.vocab.get_pad_index(),
+                bos_idx=self.vocab.get_eos_index(),
+                eos_idx=self.vocab.get_eos_index(),
+                unk_idx=self.vocab.get_unk_index(),
+            ),
+            language_vocab=ScriptVocabulary(languages),
+            max_seq_len=self.max_seq_len,
+            default_language=self.default_language,
+        )

pytext/torchscript/module.py

@@ -1,7 +1,7 @@
 #!/usr/bin/env python3
 # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
 
-from typing import List
+from typing import List, Optional
 
 import torch
 from pytext.torchscript.tensorizer.tensorizer import ScriptTensorizer
@@ -53,3 +53,50 @@ def forward(self, tokens: List[List[str]]):
         input_tensors = self.tensorizer.tensorize(tokens=squeeze_2d(tokens))
         logits = self.model(input_tensors)
         return self.output_layer(logits)
+
+
+class ScriptTokenLanguageModule(torch.jit.ScriptModule):
+    def __init__(
+        self,
+        model: torch.jit.ScriptModule,
+        output_layer: torch.jit.ScriptModule,
+        tensorizer: ScriptTensorizer,
+    ):
+        super().__init__()
+        self.model = model
+        self.output_layer = output_layer
+        self.tensorizer = tensorizer
+
+    @torch.jit.script_method
+    def forward(self, tokens: List[List[str]], languages: Optional[List[str]] = None):
+        input_tensors = self.tensorizer.tensorize(
+            tokens=squeeze_2d(tokens), languages=squeeze_1d(languages)
+        )
+        logits = self.model(input_tensors)
+        return self.output_layer(logits)
+
+
+class ScriptTokenLanguageModuleWithDenseFeature(torch.jit.ScriptModule):
+    def __init__(
+        self,
+        model: torch.jit.ScriptModule,
+        output_layer: torch.jit.ScriptModule,
+        tensorizer: ScriptTensorizer,
+    ):
+        super().__init__()
+        self.model = model
+        self.output_layer = output_layer
+        self.tensorizer = tensorizer
+
+    @torch.jit.script_method
+    def forward(
+        self,
+        tokens: List[List[str]],
+        dense_feat: List[List[float]],
+        languages: Optional[List[str]] = None,
+    ):
+        input_tensors = self.tensorizer.tensorize(
+            tokens=squeeze_2d(tokens), languages=squeeze_1d(languages)
+        )
+        logits = self.model(input_tensors, torch.tensor(dense_feat).float())
+        return self.output_layer(logits)