feat: add pipeline registry abstraction (#17905)
* feat: add pipeline registry abstraction
- added `PipelineRegistry` abstraction
- updates `add_new_pipeline.mdx` (english docs) to reflect the api addition
- migrate `check_task` and `get_supported_tasks` from
transformers/pipelines/__init__.py to
transformers/pipelines/base.py#PipelineRegistry.{check_task,get_supported_tasks}
Signed-off-by: Aaron Pham <29749331+aarnphm@users.noreply.github.com>
* fix: update with upstream/main
chore: Apply suggestions from sgugger's code review
Signed-off-by: Aaron Pham <29749331+aarnphm@users.noreply.github.com>
Co-authored-by: Sylvain Gugger <35901082+sgugger@users.noreply.github.com>
* chore: PR updates
- revert src/transformers/dependency_versions_table.py from upstream/main
- updates pipeline registry to use global variables
Signed-off-by: Aaron Pham <29749331+aarnphm@users.noreply.github.com>
* tests: add tests for pipeline registry
Signed-off-by: Aaron Pham <29749331+aarnphm@users.noreply.github.com>
* tests: add test for output warning.
Signed-off-by: Aaron Pham <29749331+aarnphm@users.noreply.github.com>
* chore: fmt and cleanup unused imports
Signed-off-by: Aaron Pham <29749331+aarnphm@users.noreply.github.com>
* fix: change imports to top of the file and address comments
Signed-off-by: Aaron Pham <29749331+aarnphm@users.noreply.github.com>
Co-authored-by: Sylvain Gugger <35901082+sgugger@users.noreply.github.com>
This commit is contained in:
@@ -111,8 +111,35 @@ of arguments for ease of use (audio files, can be filenames, URLs or pure bytes)
|
||||
|
||||
## Adding it to the list of supported tasks
|
||||
|
||||
Go to `src/transformers/pipelines/__init__.py` and fill in `SUPPORTED_TASKS` with your newly created pipeline.
|
||||
If possible it should provide a default model.
|
||||
To register your `new-task` to the list of supported tasks, provide the
|
||||
following task template:
|
||||
|
||||
```python
|
||||
my_new_task = {
|
||||
"impl": MyPipeline,
|
||||
"tf": (),
|
||||
"pt": (AutoModelForAudioClassification,) if is_torch_available() else (),
|
||||
"default": {"model": {"pt": "user/awesome_model"}},
|
||||
"type": "audio", # current support type: text, audio, image, multimodal
|
||||
}
|
||||
```
|
||||
|
||||
<Tip>
|
||||
|
||||
Take a look at the `src/transformers/pipelines/__init__.py` and the dictionary `SUPPORTED_TASKS` to see how a task is defined.
|
||||
If possible your custom task should provide a default model.
|
||||
|
||||
</Tip>
|
||||
|
||||
Then add your custom task to the list of supported tasks via
|
||||
`PIPELINE_REGISTRY.register_pipeline()`:
|
||||
|
||||
```python
|
||||
from transformers.pipelines import PIPELINE_REGISTRY
|
||||
|
||||
PIPELINE_REGISTRY.register_pipeline("new-task", my_new_task)
|
||||
```
|
||||
|
||||
|
||||
## Adding tests
|
||||
|
||||
|
||||
Reference in New Issue
Block a user