import gradio as gr
import joblib
import os

model = joblib.load('arabic_text_classifier.pkl')
vectorizer = joblib.load('tfidf_vectorizer.pkl')
label_encoder = joblib.load('label_encoder.pkl')

def predict_category(text):
    text_vector = vectorizer.transform([text])
    probabilities = model.predict_proba(text_vector)[0]
    max_prob = max(probabilities)
    predicted_category = model.predict(text_vector)[0]

    
    if max_prob < 0.5:
        return "Other"  

    
    predicted_label = label_encoder.inverse_transform([predicted_category])[0]
    return predicted_label

HF_TOKEN = os.getenv("classification")
hf_writer = gr.HuggingFaceDatasetSaver(HF_TOKEN, "crowdsourced-text-classification-data")

iface = gr.Interface(
    fn=predict_category,
    inputs=gr.Textbox(
        lines=5, 
        placeholder="Enter text in Arabic here...", 
        label="Text"  
    ),
    outputs=gr.Label(label="Predicted Category"),  
    title="Arabic Text Classification", 
    description="Enter an Arabic text to get its classification based on the trained model.",
    allow_flagging="auto",
    flagging_callback=hf_writer,
)


iface.launch(share=True)