import gradio as gr
import pickle
from huggingface_hub import hf_hub_download

model_path = hf_hub_download(repo_id="alperugurcan/toxic-model", filename="toxic_model.pkl")
vectorizer_path = hf_hub_download(repo_id="alperugurcan/toxic-model", filename="toxic_vectorizer.pkl")

model = pickle.load(open(model_path, 'rb'))
vectorizer = pickle.load(open(vectorizer_path, 'rb'))

def predict(text):
    features = vectorizer.transform([text.lower()])
    predictions = model.predict_proba(features)

    labels = ['toxic', 'severe_toxic', 'obscene', 'threat', 'insult', 'identity_hate']
    return {label: float(pred[0][1]) for label, pred in zip(labels, predictions)}

iface = gr.Interface(
    fn=predict,
    inputs="text",
    outputs="label",
    title="Toxic Comment Classifier"
)

iface.launch()