transformers/examples/research_projects/onnx/summarization

Bart + Beam Search to ONNX

Author: @fatcat-z

This folder contains an example of exporting Bart + Beam Search generation (BartForConditionalGeneration) to ONNX.

Beam Search contains a for-loop workflow, so we need to make them TorchScript-compatible for exporting to ONNX. This example shows how to make a Bart model be TorchScript-compatible by wrapping up it into a new model. In addition, some changes were made to the beam_search() function to make it TorchScript-compatible.

How to run the example

To make sure you can successfully run the latest versions of the example scripts, you have to install the library from source and install some example-specific requirements. To do this, execute the following steps in a new virtual environment:

git clone https://github.com/huggingface/transformers
cd transformers
pip install '.[onnxruntime]'

Then cd in this example folder and run

pip install -r requirements.txt

Now you can run the example command below to get the example ONNX file:

python run_onnx_exporter.py --model_name_or_path facebook/bart-base