Upload folder using huggingface_hub

pipeline.py

@@ -0,0 +1,25 @@
+from transformers import pipeline, Pipeline, AutoTokenizer, AutoConfig, AutoModelForSequenceClassification
+from transformers.pipelines import PIPELINE_REGISTRY
+import torch
+
+
+class SpanClassificationPipeline(Pipeline):
+    def __init__(self, model, tokenizer, device="cpu", **kwargs):
+        super().__init__(model=model, tokenizer=tokenizer, device=device, **kwargs)
+        self.model.to(self.device)
+        self.model.eval()
+
+    def _sanitize_parameters(self, **kwargs):
+        return {}, kwargs, {}
+
+    def preprocess(self, inputs):
+        return self.tokenizer(inputs, return_tensors="pt").to(self.device)
+
+    def _forward(self, model_inputs):
+        with torch.no_grad():
+            outputs = self.model(**model_inputs)
+        return outputs
+
+    def postprocess(self, model_outputs):
+        logits = model_outputs.logits
+        return int(torch.argmax(logits, dim=1).item())