| from transformers import pipeline, AutoModelForSequenceClassification | |
| from transformers.pipelines import PIPELINE_REGISTRY | |
| from bert_paper_classifier import SciBertPaperClassifierPipeline | |
| PIPELINE_REGISTRY.register_pipeline( | |
| "paper-classification", | |
| pipeline_class=SciBertPaperClassifierPipeline, | |
| pt_model=AutoModelForSequenceClassification, | |
| ) | |
| pipe = pipeline(task="paper-classification", model="HibiscusMaximus/scibert_paper_classification") | |
| pipe.push_to_hub("bert-paper-classification-pipeline") | |