ai-news-api / test_hf_model.py
vinitsingare's picture
feat: Add Explainable AI (XAI) and fix DistilBERT dataset bias
6197db3
Raw
History Blame Contribute Delete
1.14 kB
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
text = "Delhi restaurant fire LIVE: At least 21 people killed, several foreigners among those dead. Afire broke out at a hotel in Delhi’s Malviya Nagar on Wednesday morning (June 3, 2026), killing at least 21 people and leading to the rescue of more than 40 others."
print("Testing mrm8488/distilroberta-finetuned-fake-news...")
try:
model_name = "mrm8488/distilroberta-finetuned-fake-news"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)
inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=512)
with torch.no_grad():
logits = model(**inputs).logits
probs = torch.nn.functional.softmax(logits, dim=-1)[0]
# Check what index 0 and 1 mean for this specific model (usually 0 is reliable, 1 is fake, or vice versa)
print(f"Probabilities: {probs}")
print(f"Predicted class: {model.config.id2label[probs.argmax().item()]}")
except Exception as e:
print(f"Error with distilroberta: {e}")