huggingface / huggingface/transformers

Translation Model in ONNX: Choosable Output Formats

Open
#9,784 7 comments 3 reactions 0 assignees View on GitHub
Feature request
Dominant language
Python
Stars
166k
Forks
34.6k
Avg merge
3d 8h
Merged PRs (30d)
276

Description

# 🚀 Feature request

I am requesting to provide an option to specify the output format for the `translation_xx_to_yy` export to ONNX models. Currently, the output of [convert_graph_to_onnx.convert](https://github.com/huggingface/transformers/blob/6a346f0358a40f89ec384d441233bf54cac44f6a/src/transformers/convert_graph_to_onnx.py#L330) will provide the raw tensors as output (working prototype code under #9722)

## Motivation

When putting the models into production it would be great if one could chose, whether one wants to have the actual tensors or the output-tokens returned when exporting a translation pipeline to ONNX. Thereby, one is not forced to do a custom re-implementation of the [model.generate](https://github.com/huggingface/transformers/blob/c4d4e8bdbd25d9463d41de6398940329c89b7fb6/src/transformers/generation_utils.py#L101) function, which then uses the ONNX model instead of the torch one.

As for now, the part which is could be replaced by an ONNX inference session lives under the [model.generate](https://github.com/huggingface/transformers/blob/c4d4e8bdbd25d9463d41de6398940329c89b7fb6/src/transformers/generation_utils.py#L385) function. Using this in production would mean to keep a TranslationPipeline object with all corresponding model information and config plus an ONNX inference session.

## Your contribution

There may be multiple solutions to this problem:
1. User-specific re-implementation of model.generate (This is what Ill try to accomplish in the future)

2. Is it possible to rewrite the code under model.generate to full torch? Then it should be possible to create a custom model for all translation models, that just places this "generate layer" on top of it. I have provided an example [here](https://github.com/oborchers/sentence-transformers/blob/master/examples/onnx_inference/onnx_inference.ipynb) which adds a simple pooling layer on an already extant transformers model. (That would require more study from my side to develop a prototype and follows step 1)
3. Provide support for the [ort-customops](https://github.com/microsoft/ort-customops) library by Microsoft. Essentially, this enables ONNX to handle strings (but introduces dependency to a very experimental extension). For example, that way one can export the universal sentence encoder (including tokenizer) to ONNX. Example [here](https://github.com/onnx/tensorflow-onnx/issues/1260). I cannot provide anything useful here.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.