import gradio as gr
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
import torch

# تحميل النموذج والتوكنايزر
model_id = "methodya/arabic-summarizer-philosophy-v2"
model = AutoModelForSeq2SeqLM.from_pretrained(model_id)
tokenizer = AutoTokenizer.from_pretrained(model_id)

# نقل النموذج إلى GPU إذا كان متوفراً
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = model.to(device)
model.eval()

def summarize(text, max_length=200, num_beams=7, length_penalty=1.2):
    # تحويل القيم إلى النوع المناسب
    max_length = int(max_length)
    num_beams = int(num_beams)
    length_penalty = float(length_penalty)

    if not text.strip():
        return "الرجاء إدخال نص للتلخيص"

    # تجهيز البيانات للنموذج
    inputs = tokenizer(
        text, 
        return_tensors="pt", 
        max_length=2048,  # زيادة الحد الأقصى للمدخلات
        padding=True, 
        truncation=True
    )
    inputs = {k: v.to(device) for k, v in inputs.items()}

    # توليد الملخص مع معاملات محسنة
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_length=max_length,
            min_length=50,  # زيادة الحد الأدنى للملخص
            num_beams=num_beams,
            length_penalty=length_penalty,
            no_repeat_ngram_size=5,  # منع تكرار العبارات
            repetition_penalty=2.0,  # معاقبة التكرار
            do_sample=True,  # تفعيل العينات العشوائية
            top_k=50,  # تحديد أفضل الكلمات
            temperature=0.8,  # ضبط الإبداعية
            early_stopping=True,  # التوقف المبكر
            bad_words_ids=[[tokenizer.encode("أفهم", add_special_tokens=False)[0]], 
                         [tokenizer.encode("أعرف", add_special_tokens=False)[0]]]  # منع كلمات محددة
        )

    summary = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return summary

# تعريف واجهة Gradio
interface = gr.Interface(
    fn=summarize,
    inputs=[
        gr.Textbox(lines=8, label="النص", placeholder="أدخل النص العربي هنا..."),
        gr.Slider(100, 400, value=200, label="طول الملخص"),
        gr.Slider(1, 10, value=7, step=1, label="دقة التلخيص (num_beams)"),
        gr.Slider(0.5, 2.0, value=1.2, step=0.1, label="معامل الطول (length_penalty)")
    ],
    outputs=gr.Textbox(label="الملخص"),
    title="نظام تلخيص النصوص الفلسفية العربية",
    description="نموذج متقدم لتلخيص النصوص الفلسفية باللغة العربية",
    submit_btn="تلخيص",
    clear_btn="مسح"
)

# تشغيل التطبيق
interface.launch()