TorchScript BART models

[stephenroller]

Enable TorchScripting of a BART-based TransformerGeneratorModel. Support for other Transformer variants may come at a later time.

Distilled MultiWoZ-tuned model (BART)

Known good conversation

Call to interact with model:

parlai interactive \
--model-file ${MODEL_FILE} \
--model bart \
--no-cuda

Output:

Enter Your Message: I am looking for a restaurant in the west part of town.
[Bart]: APICALL: Restaurant.area = west
Enter Your Message: APIRESP: Restaurant 14 matches
[Bart]: there are [value_count] restaurant -s in the [value_area] part of town . do you have a preference ?

Exporting model

Call to export model:

parlai torchscript \
--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'

Output:

Generating given the original unscripted module:
 TEXT: I am looking for a restaurant in the west part of town.
LABEL: APICALL: Restaurant.area = west
 TEXT: APIRESP: Restaurant 14 matches
LABEL: there are [value_count] restaurant -s in the [value_area] part of town . do you have a preference ?

Generating given the scripted module:
 TEXT: I am looking for a restaurant in the west part of town.
LABEL: APICALL: Restaurant.area = west
 TEXT: APIRESP: Restaurant 14 matches
LABEL: there are [value_count] restaurant -s in the [value_area] part of town . do you have a preference ?

Testing loading in and using exported model

Call:

python parlai/torchscript/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'

Output:

Generating given the scripted module:
 TEXT: I am looking for a restaurant in the west part of town.
LABEL: APICALL: Restaurant.area = west
 TEXT: APIRESP: Restaurant 14 matches
LABEL: there are [value_count] restaurant -s in the [value_area] part of town . do you have a preference ?

Additional testing

• Unit test: pytest tests/test_torchscript.py
• Agent: parlai interactive --model-file ~/_test_scripted_model__bart.pt --model parlai.torchscript.agents:TorchScriptAgent

[EricMichaelSmith]

Ha, reminds me of the shenanigans I had to go through with ONNX (i.e. removing FusedLayerNorm, etc.)

[stephenroller]

Ha, reminds me of the shenanigans I had to go through with ONNX (i.e. removing FusedLayerNorm, etc.)

ya that's what we need to discuss lol

[EricMichaelSmith]

@stephenroller Yes, I've re-split the incremental-state reordering logic back into TransformerDecoderLayer and MultiHeadAttention