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

# =================================================================
# ==     ⚠️ قم بتغيير هذا السطر فقط ⚠️     ==
# =================================================================
# ضع هنا الاسم الكامل للمستودع الذي رفعت إليه نموذجك
MODEL_ID = "FaizTech/my-tinyllama-test"  # <--- مثال، غيره إلى اسم نموذجك
# =================================================================


# --- لا داعي لتعديل أي شيء بعد هذا السطر ---

# رسالة لتعقب حالة التحميل في الـ Logs
print(f"بدء تحميل النموذج: {MODEL_ID}")

# تحميل الـ Tokenizer أولاً
try:
    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
except Exception as e:
    print(f"خطأ في تحميل الـ Tokenizer: {e}")
    raise

# إنشاء الـ pipeline الذي سيقوم بتحميل النموذج وتشغيله
try:
    pipe = pipeline(
        "text-generation",
        model=MODEL_ID,
        tokenizer=tokenizer,
        torch_dtype=torch.bfloat16,
        device_map="auto" # سيستخدم الـ CPU أو الـ GPU المتاح تلقائيًا
    )
    print("تم تحميل النموذج بنجاح!")
except Exception as e:
    print(f"خطأ في تحميل النموذج: {e}")
    raise


# الدالة الرئيسية التي ستعالج طلبات الـ API والواجهة
def get_response(prompt_text):
    """
    هذه الدالة تأخذ نصًا وتولد ردًا من النموذج.
    """
    # استخدام قالب المحادثة الجاهز في الـ tokenizer
    messages = [{"role": "user", "content": prompt_text}]
    prompt = pipe.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)

    # توليد الرد
    outputs = pipe(
        prompt,
        max_new_tokens=256,
        do_sample=True,
        temperature=0.7,
        top_p=0.95
    )
    
    # استخراج النص المولد فقط
    response = outputs[0]["generated_text"]
    
    # تنظيف الرد من النص الأصلي (اختياري، لتحسين العرض)
    # هذا السطر يزيل السؤال ويعرض الإجابة فقط
    cleaned_response = response.replace(prompt, "").strip()
    
    return cleaned_response

# إنشاء واجهة Gradio البسيطة
demo = gr.Interface(
    fn=get_response,
    inputs=gr.Textbox(label="أدخل سؤالك أو أمرك هنا", lines=4),
    outputs=gr.Textbox(label="إجابة النموذج", lines=4),
    title=f"واجهة لنموذج {MODEL_ID.split('/')[-1]}",
    description="هذا تطبيق ويب بسيط لواجهة API لنموذجي الخاص."
)

# تشغيل التطبيق وإتاحة الـ API
demo.queue().launch()