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

# اسم مستودع الموديل متاعك
model_name = "Dav66/Te"

# تحميل التوكنيزر
tokenizer = AutoTokenizer.from_pretrained(model_name)

# تحميل الموديل مع تفعيل offloading على القرص وتحديد نوع البيانات
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    device_map="auto",               # توزيع تلقائي على CPU
    offload_folder="./offload",     # مجلد لتخزين الأجزاء المؤقتة
    torch_dtype=torch.float16       # تقليل حجم الأوزان
)

# دالة التوليد (باستهلاك ذاكرة منخفض)
def generate_text(prompt):
    with torch.no_grad():  # تقليل استهلاك الذاكرة
        inputs = tokenizer(prompt, return_tensors="pt")
        outputs = model.generate(**inputs, max_new_tokens=50)  # توليد قصير نسبياً
        text = tokenizer.decode(outputs[0], skip_special_tokens=True)
        return text

# واجهة Gradio
iface = gr.Interface(
    fn=generate_text,
    inputs=gr.Textbox(lines=2, placeholder="اكتب هنا النص..."),
    outputs="text",
    title="توليد نص بالعربي",
    description="تطبيق بسيط يعرض موديلك على Hugging Face Space بكفاءة على CPU"
)

iface.launch()