import gradio as gr
from transformers import pipeline

# تحميل النموذج المدرب أو موديل جاهز
classifier = pipeline(
    "text-classification",
    model="aubmindlab/bert-base-arabertv2",
    tokenizer="aubmindlab/bert-base-arabertv2",
    top_k=None,  # يرجع لك كل الاحتمالات
)

# دالة التصنيف
def classify_dialect(text):
    outputs = classifier(text)
    # ترتيب النتائج حسب الاحتمالية
    outputs = sorted(outputs, key=lambda x: x['score'], reverse=True)
    result = {item['label']: round(item['score'], 3) for item in outputs}
    return result

# واجهة Gradio
demo = gr.Interface(
    fn=classify_dialect,
    inputs=gr.Textbox(label="أدخل النص باللهجة العربية"),
    outputs=gr.Label(label="توقع اللهجة"),
    title="تصنيف اللهجات العربية",
    description="أدخل نص عربي وسنحاول تصنيف اللهجة باستخدام نموذج BERT مدرب مسبقاً.",
)

# تشغيل التطبيق
if __name__ == "__main__":
    demo.launch()