pipeline registry?

pipeline.py

@@ -1,15 +1,10 @@
-from typing import List, Dict
-import json
+from transformers import PIPELINE_REGISTRY, TokenClassificationPipeline
 import torch
-from transformers import BertTokenizerFast, BertForTokenClassification
-
-class BiasNERPipeline:
-    def __init__(self, model_path: str = 'maximuspowers/bias-detection-ner'):
-        self.tokenizer = BertTokenizerFast.from_pretrained('bert-base-uncased')
-        self.model = BertForTokenClassification.from_pretrained(model_path)
-        self.model.eval()
-        self.model.to('cuda' if torch.cuda.is_available() else 'cpu')
 
+@PIPELINE_REGISTRY.register_pipeline(task="multi_label_token_classification", pipeline_class=None)
+class MultiLabelTokenClassificationPipeline(TokenClassificationPipeline):
+    def __init__(self, model, tokenizer, **kwargs):
+        super().__init__(model=model, tokenizer=tokenizer, **kwargs)
         self.id2label = {
             0: 'O',
             1: 'B-STEREO',
@@ -20,23 +15,16 @@ class BiasNERPipeline:
             6: 'I-UNFAIR'
         }
 
-    def __call__(self, inputs: str) -> str:
-        tokenized_inputs = self.tokenizer(inputs, return_tensors="pt", padding=True, truncation=True, max_length=128)
-        input_ids = tokenized_inputs['input_ids'].to(self.model.device)
-        attention_mask = tokenized_inputs['attention_mask'].to(self.model.device)
-
-        with torch.no_grad():
-            outputs = self.model(input_ids=input_ids, attention_mask=attention_mask)
-            logits = outputs.logits
+    def postprocess(self, model_outputs, **kwargs):
+        results = []
+        for logits, tokens in zip(model_outputs[0], model_outputs[1]):
             probabilities = torch.sigmoid(logits)
             predicted_labels = (probabilities > 0.5).int()
-
-        result = []
-        tokens = self.tokenizer.convert_ids_to_tokens(input_ids[0])
-        for i, token in enumerate(tokens):
-            if token not in self.tokenizer.all_special_tokens:
-                label_indices = (predicted_labels[0][i] == 1).nonzero(as_tuple=False).squeeze(-1)
-                labels = [self.id2label[idx.item()] for idx in label_indices] if label_indices.numel() > 0 else ['O']
-                result.append({"token": token, "labels": labels})
-
-        return json.dumps(result, indent=4)
+            token_results = []
+            for i, token in enumerate(tokens):
+                if token not in self.tokenizer.all_special_tokens:
+                    label_indices = (predicted_labels[i] == 1).nonzero(as_tuple=False).squeeze(-1)
+                    labels = [self.id2label[idx.item()] for idx in label_indices] if label_indices.numel() > 0 else ['O']
+                    token_results.append({"token": token, "labels": labels})
+            results.append(token_results)
+        return results