# 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")
