Update pipeline.py

pipeline.py

@@ -1,11 +1,8 @@
-from typing import Dict, List
+from typing import List, Dict
 import json
 import torch
-import numpy as np
 from transformers import BertTokenizerFast, BertForTokenClassification
 
-# this is so that we can use a custom pipeline (mostly parsing outputs) with the pipeline module
-
 class BiasNERPipeline:
     def __init__(self, model_path: str = 'maximuspowers/bias-detection-ner'):
         self.tokenizer = BertTokenizerFast.from_pretrained('bert-base-uncased')
@@ -13,7 +10,6 @@ class BiasNERPipeline:
         self.model.eval()
         self.model.to('cuda' if torch.cuda.is_available() else 'cpu')
 
-        # label mapping
         self.id2label = {
             0: 'O',
             1: 'B-STEREO',
@@ -24,20 +20,17 @@ class BiasNERPipeline:
             6: 'I-UNFAIR'
         }
 
-    def __call__(self, inputs: str) -> List[Dict[str, float]]:
-        # tokenize
+    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)
 
-        # run model
         with torch.no_grad():
             outputs = self.model(input_ids=input_ids, attention_mask=attention_mask)
             logits = outputs.logits
             probabilities = torch.sigmoid(logits)
             predicted_labels = (probabilities > 0.5).int()
 
-        # format output
         result = []
         tokens = self.tokenizer.convert_ids_to_tokens(input_ids[0])
         for i, token in enumerate(tokens):
@@ -46,4 +39,4 @@ class BiasNERPipeline:
                 labels = [self.id2label[idx.item()] for idx in label_indices] if label_indices.numel() > 0 else ['O']
                 result.append({"token": token, "labels": labels})
 
-        return result
+        return json.dumps(result, indent=4)