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")