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)