newsmediabias/UnBIAS-RoBERTa-NER

# Load model directly
from transformers import AutoTokenizer, AutoModelForTokenClassification

tokenizer = AutoTokenizer.from_pretrained("newsmediabias/UnBIAS-Roberta-NER")
model = AutoModelForTokenClassification.from_pretrained("newsmediabias/UnBIAS-Roberta-NER")

# Example batch of sentences
sentences = [
    "The corrupt politician embezzled funds.",
    "Immigrants are causing a surge in crime.",
    "The movie star is an idiot for their political views.",
    "Only a fool would believe in climate change.",
    "The new policy will destroy the economy."
]

# Tokenize the batch
encoding = tokenizer(sentences, return_tensors='pt', padding=True, truncation=True)

# Get model predictions
outputs = model(**encoding)

# Apply softmax to the output logits to get probabilities
predictions = torch.nn.functional.softmax(outputs.logits, dim=-1)

# Get the highest probability labels for each token
predicted_labels = torch.argmax(predictions, dim=-1)

# Define a mapping for the labels
label_mapping = {
    0: "O",      # No bias
    1: "B-BIAS", # Beginning of a biased sequence
    2: "I-BIAS"  # Inside a biased sequence
}

# Convert predicted labels to their corresponding label names using the mapping
labels = [[label_mapping[label_id.item()] for label_id in sentence_labels] for sentence_labels in predicted_labels]

# Align labels with the words in the sentences
aligned_labels = []
for i, sentence_labels in enumerate(labels):
    # Get the tokens from the original sentence
    tokens = tokenizer.convert_ids_to_tokens(encoding['input_ids'][i])
    # Only consider labels for tokens that are not special tokens
    sentence_labels = [label for token, label in zip(tokens, sentence_labels) if token not in tokenizer.all_special_tokens]
    aligned_labels.append(sentence_labels)

# Print the aligned labels for each sentence
for sentence, labels in zip(sentences, aligned_labels):
    print(f"Sentence: {sentence}\nLabels: {labels}\n")