import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
import os
from huggingface_hub import login

# تسجيل الدخول
login(token=os.environ.get('HUGGING_FACE_HUB_TOKEN'))

# تهيئة النموذج
model_name = "google/gemma-2b-it"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

device = 'cuda' if torch.cuda.is_available() else 'cpu'
model = model.to(device)

def generate_summary(text):
    prompt = f"""لخص هذا النص بدقة:
{text}

التلخيص يجب أن يتضمن:
- فقرة قصيرة تشرح الفكرة الأساسية
- ثلاث نقاط رئيسية
"""
    
    inputs = tokenizer(prompt, return_tensors="pt").to(device)
    outputs = model.generate(
        **inputs,
        max_length=256,  # تقليل الطول لتسريع الاستجابة
        temperature=0.3  # تقليل العشوائية للحصول على نتائج أكثر دقة
    )
    
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

css = ".rtl-text { direction: rtl; text-align: right; }"

interface = gr.Interface(
    fn=generate_summary,
    inputs=gr.Textbox(label="النص", lines=6, elem_classes="rtl-text"),
    outputs=gr.Textbox(label="الملخص", lines=6, elem_classes="rtl-text"),
    title="ملخص النصوص",
    theme=gr.themes.Soft(),
    css=css
)

interface.launch()