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. |