import gradio as gr
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch

# Load model and tokenizer
model = AutoModelForSequenceClassification.from_pretrained("remomo/URLClassifier")
tokenizer = AutoTokenizer.from_pretrained("remomo/URLClassifier")

# Define prediction function
def classify_url(url):
    inputs = tokenizer(url, return_tensors="pt", truncation=True, padding=True)
    with torch.no_grad():
        outputs = model(**inputs)
        pred = torch.argmax(outputs.logits, dim=1).item()
        return "🔒 Safe" if pred == 0 else "⚠️ Unsafe"

# Create Gradio interface
iface = gr.Interface(fn=classify_url, 
                     inputs=gr.Textbox(label="Enter URL"), 
                     outputs=gr.Textbox(label="Classification"),
                     title="URL Safety Checker")

iface.launch(share=True)