File size: 1,285 Bytes
60750b5 a2af29b 60750b5 | 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 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 | import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import gradio as gr
# Load model and tokenizer
model_name = "sansanmaw/ssm_model"
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)
model.to(device)
model.eval()
# Define ID-to-label mapping
id2lbl = {0: "fake", 1: "true"}
# Prediction function
def predict_news(text):
inputs = tokenizer(
text,
padding=True,
truncation=True,
max_length=512,
return_tensors="pt"
)
inputs = {k: v.to(device) for k, v in inputs.items()}
with torch.no_grad():
outputs = model(**inputs)
logits = outputs.logits
predictions = torch.argmax(logits, dim=-1).cpu().numpy()
predicted_label = id2lbl[predictions[0]]
return predicted_label
# Gradio interface
interface = gr.Interface(
fn=predict_news,
inputs=gr.Textbox(lines=8, placeholder="Paste news article here..."),
outputs="text",
title="Fake News Detector (Roberta)",
description="Paste a news article to check if it's real or fake."
)
# Launch app
if __name__ == "__main__":
interface.launch()
|