File size: 593 Bytes
5fa1a76 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 |
If we have saved this in
a file named pair_classification.py, we can then import it and register it like this:
from pair_classification import PairClassificationPipeline
from transformers.pipelines import PIPELINE_REGISTRY
from transformers import AutoModelForSequenceClassification, TFAutoModelForSequenceClassification
PIPELINE_REGISTRY.register_pipeline(
"pair-classification",
pipeline_class=PairClassificationPipeline,
pt_model=AutoModelForSequenceClassification,
tf_model=TFAutoModelForSequenceClassification,
)
Once this is done, we can use it with a pretrained model. |