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()