ZetangForward commited on
Commit
eb143e0
·
verified ·
1 Parent(s): 5d4a357

Upload folder using huggingface_hub

Browse files
Files changed (1) hide show
  1. pipeline.py +8 -1
pipeline.py CHANGED
@@ -22,4 +22,11 @@ class SpanClassificationPipeline(Pipeline):
22
 
23
  def postprocess(self, model_outputs):
24
  logits = model_outputs.logits
25
- return int(torch.argmax(logits, dim=1).item())
 
 
 
 
 
 
 
 
22
 
23
  def postprocess(self, model_outputs):
24
  logits = model_outputs.logits
25
+ return int(torch.argmax(logits, dim=1).item())
26
+
27
+
28
+ PIPELINE_REGISTRY.register_pipeline(
29
+ task="spancnn-classification",
30
+ pipeline_class=SpanClassificationPipeline,
31
+ pt_model=AutoModelForSequenceClassification,
32
+ )