import gradio as gr
import joblib
import re
from nltk.corpus import stopwords
import nltk

nltk.download('stopwords', quiet=True)
stop_words = set(stopwords.words('english'))

svm_model = joblib.load('svm_model.pkl')
vectorizer = joblib.load('tfidf_vectorizer.pkl')

label_map = {
    1: 'World',
    2: 'Sports',
    3: 'Business',
    4: 'Sci/Tech'
}

def clean_text(text):
    text = text.lower()
    text = re.sub(r'[^a-zA-Z]', ' ', text)
    words = text.split()
    words = [w for w in words if w not in stop_words]
    return " ".join(words)

def predict(text):
    if not text.strip():
        return "Please enter some text."
    cleaned = clean_text(text)
    vectorized = vectorizer.transform([cleaned])
    prediction = svm_model.predict(vectorized)[0]
    return label_map[prediction]

demo = gr.Interface(
    fn=predict,
    inputs=gr.Textbox(
        lines=4,
        placeholder="Paste a news headline or article here...",
        label="News Text"
    ),
    outputs=gr.Textbox(label="Predicted Category"),
    title="AG News Classifier",
    description="Classifies news articles into: World, Sports, Business, or Sci/Tech. Built with SVM + TF-IDF trained on 120,000 articles.",
    examples=[
        ["NASA discovers water on Mars surface"],
        ["Stock markets crash amid banking crisis"],
        ["Ronaldo scores hat-trick in World Cup qualifier"],
        ["World leaders gather for climate summit in Paris"],
        ["Apple launches new AI chip for iPhones"]
    ],
    theme=gr.themes.Soft()
)

demo.launch()