Upload folder using huggingface_hub

config.json

@@ -5,6 +5,15 @@
   ],
   "attention_probs_dropout_prob": 0.1,
   "classifier_dropout": null,
+  "custom_pipelines": {
+    "newsagency-ner": {
+      "impl": "newsagency_ner.NewsAgencyModelPipeline",
+      "pt": [
+        "AutoModelForTokenClassification"
+      ],
+      "tf": []
+    }
+  },
   "hidden_act": "gelu",
   "hidden_dropout_prob": 0.1,
   "hidden_size": 768,

newsagency_ner.py

@@ -167,8 +167,11 @@ class NewsAgencyModelPipeline(Pipeline):
         input_ids = torch.tensor([inputs["input_ids"]], dtype=torch.long).to(
             self.model.device
         )
+        attention_mask = torch.tensor([inputs["attention_mask"]], dtype=torch.long).to(
+            self.model.device
+        )
         with torch.no_grad():
-            outputs = self.model(input_ids)
+            outputs = self.model(input_ids, attention_mask)
         return outputs, text_sentence
 
     def postprocess(self, outputs, **kwargs):