TorchScript BART models

parlai/agents/bart/modules.py

@@ -7,7 +7,7 @@
 """
 import torch
 import torch.nn.functional as F
-from typing import Any, Dict, Union, List, Optional
+from typing import Dict
 from parlai.agents.transformer.modules import TransformerGeneratorModel
 
 
@@ -50,10 +50,8 @@ def _get_initial_forced_decoder_input(self, bsz: int, inputs: torch.LongTensor):
         return torch.cat([tens, inputs], 1)
 
     def reorder_decoder_incremental_state(
-        self,
-        incremental_state: Dict[str, Any],
-        inds: Union[List[int], torch.LongTensor],
-    ) -> Optional[Dict[str, Any]]:
+        self, incremental_state: Dict[str, torch.Tensor], inds: torch.LongTensor
+    ) -> Dict[str, torch.Tensor]:
         """
         Incremental state is weird to handle when we seed decoder with two inputs
         initially.
@@ -62,14 +60,12 @@ def reorder_decoder_incremental_state(
         assert incremental_state is not None
         assert len(incremental_state) > 0
 
-        for incr_state_l in incremental_state.values():
-            assert 'self_attn' in incr_state_l
-            assert 'prev_mask' in incr_state_l['self_attn']
-            self_attn_mask = incr_state_l['self_attn']['prev_mask']
-            # check this is on the very first run with incremental state
-            if self_attn_mask.ndim == 3 and tuple(self_attn_mask.shape[1:]) == (2, 2):
-                # cut off the inappropriate incremental state
-                incr_state_l['self_attn']['prev_mask'] = self_attn_mask[:, -1:, :]
+        assert 'self_attn_prev_mask' in incremental_state
+        self_attn_mask = incremental_state['self_attn_prev_mask']
+        # check this is on the very first run with incremental state
+        if self_attn_mask.ndim == 4 and tuple(self_attn_mask.shape[1:3]) == (2, 2):
+            # cut off the inappropriate incremental state
+            incremental_state['self_attn_prev_mask'] = self_attn_mask[:, -1:, :, :]
 
         return super().reorder_decoder_incremental_state(incremental_state, inds)
 

parlai/agents/jit/README.md

@@ -0,0 +1,25 @@
+# Agent exported to TorchScript (JIT compilation)
+
+This agent will read in a ParlAI agent that has been exported to TorchScript with JIT compilation, for use in greedy-search inference on CPU. This allows inference to be run on models without using any ParlAI overhead, either for tokenization or for the forward passes through the model. Currently, only BART models are supported.
+
+Sample call for exporting a BART model to TorchScript:
+```
+parlai jit_export \
+--model-file ${MODEL_FILE} \
+--model bart \
+--no-cuda \
+--scripted-model-file ~/_test_scripted_model__bart.pt \
+--input 'I am looking for a restaurant in the west part of town.|APIRESP: Restaurant 14 matches'
+```
+
+Interacting with an exported model using `parlai interactive`:
+```
+parlai interactive --model-file ~/_test_scripted_model__bart.pt --model jit
+```
+
+Loading in and running inference on an exported model, without any ParlAI overhead:
+```
+python parlai/agents/jit/scripts/test_exported_model.py \
+--scripted-model-file ~/_test_scripted_model__bart.pt \
+--input 'I am looking for a restaurant in the west part of town.|APIRESP: Restaurant 14 matches'
+```

parlai/agents/jit/__init__.py

@@ -0,0 +1,5 @@
+#!/usr/bin/env python3
+
+# Copyright (c) Facebook, Inc. and its affiliates.
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.

parlai/agents/jit/jit.py

@@ -0,0 +1,70 @@
+#!/usr/bin/env python3
+
+# Copyright (c) Facebook, Inc. and its affiliates.
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+from typing import List
+
+import torch
+
+from parlai.core.agents import Agent
+from parlai.core.message import Message
+from parlai.core.opt import Opt
+from parlai.utils.io import PathManager
+
+
+class JitAgent(Agent):
+    """
+    ParlAI agent exported to TorchScript with JIT compilation and then loaded from disk.
+
+    Metrics and batch act are currently unsupported, and CUDA is unsupported because
+    TorchScripting is currently CPU-only.
+    """
+
+    def __init__(self, opt: Opt, shared=None):
+
+        super().__init__(opt=opt, shared=shared)
+        with PathManager.open(self.opt['model_file'], "rb") as f:
+            self.module = torch.jit.load(f)
+
+        # Track incoming history strings
+        self.history: List[str] = []
+
+    def share(self):
+        """
+        Share the scripted module object.
+        """
+        shared = super().share()
+        shared['module'] = self.module
+        return shared
+
+    def observe(self, observation: Message) -> Message:
+        # TODO: support self._validate_observe_invariants() method of TorchAgent
+
+        self.history.append(observation['text'])
+
+        return super().observe(observation)
+
+    def self_observe(self, self_message: Message) -> None:
+        # TODO: support self._validate_self_observe_invariants() method of TorchAgent
+
+        assert self.observation is not None
+        if self.observation['episode_done']:
+            # oh this was the last example in the episode. reset the history
+            self.history = []
+            # additionally mark the last observation as invalid
+            self.observation = None
+        else:
+            self.history.append(self_message['text'])
+
+    def reset(self):
+        super().reset()
+        self.history = []
+
+    def act(self) -> Message:
+        response_text = self.module('\n'.join(self.history))
+        response = Message({'text': response_text, 'episode_done': False})
+        # self.observation will determine if we're going onto a new episode
+        self.self_observe(response)
+        return response

parlai/agents/jit/scripts/__init__.py

@@ -0,0 +1,5 @@
+#!/usr/bin/env python3
+
+# Copyright (c) Facebook, Inc. and its affiliates.
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.

parlai/agents/jit/scripts/test_exported_model.py

@@ -0,0 +1,48 @@
+#!/usr/bin/env python3
+
+# Copyright (c) Facebook, Inc. and its affiliates.
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import argparse
+from typing import List
+
+import torch.jit
+
+from parlai.utils.io import PathManager
+
+
+def test_exported_model(scripted_model_file: str, inputs: List[str]):
+
+    with PathManager.open(scripted_model_file, "rb") as f:
+        scripted_module = torch.jit.load(f)
+
+    print('\nGenerating given the scripted module:')
+    context = []
+    for input_ in inputs:
+        print(' TEXT: ' + input_)
+        context.append(input_)
+        label = scripted_module('\n'.join(context))
+        print("LABEL: " + label)
+        context.append(label)
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        '-smf',
+        '--scripted-model-file',
+        type=str,
+        help='Where to load the scripted model checkpoint from',
+    )
+    parser.add_argument(
+        "-i",
+        "--input",
+        type=str,
+        default="hello world",
+        help="Test input string to pass into the encoder of the scripted model. Separate lines with a pipe",
+    )
+    args = parser.parse_args()
+    test_exported_model(
+        scripted_model_file=args.scripted_model_file, inputs=args.input.split('|')
+    )