File size: 1,079 Bytes
4dfb4e3 d4515d7 4dfb4e3 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 |
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
model_path = "./training/bert-allsides-bias-detector/checkpoint-10494"
tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModelForSequenceClassification.from_pretrained(model_path)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
label_mapping = {0: "Left", 1: "Center", 2: "Right"}
def predict_bias(text):
"""Predicts the political bias of the given text."""
inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=512).to(device)
with torch.no_grad():
outputs = model(**inputs)
logits = outputs.logits
predicted_class = torch.argmax(logits, dim=-1).item()
return label_mapping[predicted_class]
if __name__ == "__main__":
while True:
text = input("\nEnter text to classify (or type 'exit' to quit): ")
if text.lower() == "exit":
break
bias_label = predict_bias(text)
print(f"Predicted Bias: {bias_label}")
|