Add Flax Models to Pipelines

[ncoop57]

🚀 Feature request

Hi y'all, I am trying a GPTNeo Flax model and am wanting to use it in the text generation pipeline. However, it is currently not supported. From looking at the current implementation of flax models and the text generation pipeline, it should be a relatively easy (famous last words) addition.

Motivation

HF is heavily integrating Flax models (which I think is awesome!) into the library and has parroted many of the already existing parts of the transformers library to flax models, similar to what was done for TF models. The addition of the support of Flax models in the pipeline API will help those who are working with pure Flax models, especially for applications that will be use the model to accomplish some task.

Your contribution

I would be willing to open a PR if one is not currently underway (I looked for one and didn't find any). However, I am new to flax so if the task is more difficult than I expect I probably will be not able to complete it.