| from typing import List, Dict | |
| import json | |
| 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') | |
| self.id2label = { | |
| 0: 'O', | |
| 1: 'B-STEREO', | |
| 2: 'I-STEREO', | |
| 3: 'B-GEN', | |
| 4: 'I-GEN', | |
| 5: 'B-UNFAIR', | |
| 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 | |
| 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) | |