import gradio as gr
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
import torch.nn.functional as F

# تحميل نموذج جاهز (غيّري الاسم إلى موديلك الخاص إذا عندك)
MODEL_NAME = "aubmindlab/bert-base-arabertv2"

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME)
model.eval()

# تعريف أسماء المستويات حسب الفئات (LABEL_0, LABEL_1, LABEL_2)
# غيّري الترتيب لو عندك موديل مدرَّب بترتيب مختلف
id2label = {
    0: "سهل",
    1: "متوسط",
    2: "صعب"
}

# دالة التنبؤ بالنص
def analyze_readability(text):
    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=512)
    with torch.no_grad():
        outputs = model(**inputs)
        probs = F.softmax(outputs.logits, dim=-1)
        pred = int(torch.argmax(probs))
        confidence = float(probs[0][pred])
    
    label = id2label.get(pred, "غير معروف")
    return f"تصنيف هذه النشرة: {label}"

# واجهة Gradio
iface = gr.Interface(
    fn=analyze_readability,
    inputs=gr.Textbox(label="أدخل نصًا بالعربية"),
    outputs="text",
    title="تحليل قابلية القراءة للنصوص الطبية",
    description="أدخل فقرة من النشرة الطبية ليقوم النظام بتقدير مستوى القراءة: سهل، متوسط، أو صعب."
)

iface.launch()