use ScriptXLMTensorizer (#1123)

Summary:
Pull Request resolved: #1123

use ScriptXLMTensorizer

Differential Revision: D18364254

fbshipit-source-id: 20966a39aa3631cd84cfb9a778bd19c7f8d03cc8

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):
@@ -137,3 +139,23 @@ 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=None, default_language="en"):
+        if languages is None:
+            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=default_language,
+        )