|
It's also relatively common to support many different types |
|
of arguments for ease of use (audio files, which can be filenames, URLs or pure bytes) |
|
Adding it to the list of supported tasks |
|
To register your new-task to the list of supported tasks, you have to add it to the PIPELINE_REGISTRY: |
|
thon |
|
from transformers.pipelines import PIPELINE_REGISTRY |
|
PIPELINE_REGISTRY.register_pipeline( |
|
"new-task", |
|
pipeline_class=MyPipeline, |
|
pt_model=AutoModelForSequenceClassification, |
|
) |
|
|
|
You can specify a default model if you want, in which case it should come with a specific revision (which can be the name of a branch or a commit hash, here we took "abcdef") as well as the type: |
|
python |
|
PIPELINE_REGISTRY.register_pipeline( |
|
"new-task", |
|
pipeline_class=MyPipeline, |
|
pt_model=AutoModelForSequenceClassification, |
|
default={"pt": ("user/awesome_model", "abcdef")}, |
|
type="text", # current support type: text, audio, image, multimodal |
|
) |
|
Share your pipeline on the Hub |
|
To share your custom pipeline on the Hub, you just have to save the custom code of your Pipeline subclass in a |
|
python file. |