commit files to HF hub

tunBertClassificationPipeline.py

@@ -1,7 +1,11 @@
-from transformers import Pipeline
+from transformers import Pipeline, AutoModelForSequenceClassification,AutoTokenizer
 import torch
+from transformers.pipelines import PIPELINE_REGISTRY
 
 class TBCP(Pipeline):
+    def __init__(self,**kwargs):
+      Pipeline.__init__(self,**kwargs)
+      self.tokenizer = AutoTokenizer.from_pretrained(kwargs["tokenizer"])
     def _sanitize_parameters(self, **kwargs):
         postprocess_kwargs = {}
         if "text_pair" in kwargs:
@@ -19,7 +23,17 @@ class TBCP(Pipeline):
         probabilities = torch.nn.functional.softmax(logits, dim=-1)
         
         best_class = probabilities.argmax().item()
-        label = self.model.config.id2label[best_class]
-        score = probabilities.squeeze()[best_class].item()
+        label = f"Label_{best_class}"
+        # score = probabilities.squeeze()[best_class].item()
         logits = logits.squeeze().tolist()
-        return {"label": label, "score": score, "logits": logits}
+        return {"label": label, 
+                # "score": score, 
+                "logits": logits}
+
+PIPELINE_REGISTRY.register_pipeline(
+    "TunBERT-classifier",
+    pipeline_class=TBCP,
+    pt_model=AutoModelForSequenceClassification,
+    default={"pt": ("not-lain/TunBERT", "main")},
+    type="text",  # current support type: text, audio, image, multimodal
+)