Skip to content

Commit

Permalink
[docs] Follow up register_pipeline (#35310)
Browse files Browse the repository at this point in the history
example json
  • Loading branch information
stevhliu authored Dec 20, 2024
1 parent 94fe0b9 commit 608e163
Showing 1 changed file with 17 additions and 1 deletion.
18 changes: 17 additions & 1 deletion docs/source/en/add_new_pipeline.md
Original file line number Diff line number Diff line change
Expand Up @@ -184,7 +184,7 @@ class PairClassificationPipeline(Pipeline):
```

The implementation is framework agnostic, and will work for PyTorch and TensorFlow models. If we have saved this in
a file named `pair_classification.py`, we can then import it and register it like this. The [register_pipeline](https://github.com/huggingface/transformers/blob/9feae5fb0164e89d4998e5776897c16f7330d3df/src/transformers/pipelines/base.py#L1387) function registers the pipeline details (task type, pipeline class, supported backends) to a models `config.json` file.
a file named `pair_classification.py`, we can then import it and register it like this.

```py
from pair_classification import PairClassificationPipeline
Expand All @@ -199,6 +199,22 @@ PIPELINE_REGISTRY.register_pipeline(
)
```

The [register_pipeline](https://github.com/huggingface/transformers/blob/9feae5fb0164e89d4998e5776897c16f7330d3df/src/transformers/pipelines/base.py#L1387) function registers the pipeline details (task type, pipeline class, supported backends) to a models `config.json` file.

```json
"custom_pipelines": {
"pair-classification": {
"impl": "pair_classification.PairClassificationPipeline",
"pt": [
"AutoModelForSequenceClassification"
],
"tf": [
"TFAutoModelForSequenceClassification"
],
}
},
```

Once this is done, we can use it with a pretrained model. For instance `sgugger/finetuned-bert-mrpc` has been
fine-tuned on the MRPC dataset, which classifies pairs of sentences as paraphrases or not.

Expand Down

0 comments on commit 608e163

Please sign in to comment.