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.
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 .
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