# 1. تحديث وتثبيت المكتبات لضمان التوافق
!pip install -q -U transformers accelerate bitsandbytes gradio

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

# 2. إعداد النموذج (Qwen 2.5 - 1.5B)
model_name = "Qwen/Qwen2.5-1.5B-Instruct"

# إعدادات الـ 4-bit لتقليل استهلاك الذاكرة
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.float16
)

print("⏳ جاري تشغيل Zeus... الرجاء الانتظار قليلاً...")

tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    quantization_config=bnb_config,
    device_map="auto"
)

# 3. شخصية Zeus
system_prompt = """أنت Zeus، مساعد ذكي وودود. تتحدث العربية بطلاقة وتفهم اللهجات. 
أنت خبير تقني وتقدم إجابات منظمة وواضحة جداً بالنقاط."""

# 4. دالة المعالجة (تعديل طفيف لتتوافق مع نظام الرسائل الجديد)
def chat_with_zeus(message, history):
    # تحويل التاريخ إلى تنسيق Qwen
    full_history = [{"role": "system", "content": system_prompt}]
    for h in history:
        full_history.append({"role": "user", "content": h[0]})
        full_history.append({"role": "assistant", "content": h[1]})
    
    full_history.append({"role": "user", "content": message})
    
    # تحويل الرسائل لنص يفهمه النموذج
    text = tokenizer.apply_chat_template(full_history, tokenize=False, add_generation_prompt=True)
    model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
    
    # التوليد
    generated_ids = model.generate(**model_inputs, max_new_tokens=512, temperature=0.7, do_sample=True)
    
    # استخراج الرد
    response_ids = [output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)]
    return tokenizer.batch_decode(response_ids, skip_special_tokens=True)[0]

# 5. الواجهة الرسومية (الإصدار المتوافق والمبسط)
# قمنا بإزالة الوسائط التي تسبب أخطاء واستخدمنا الوضع القياسي المستقر
demo = gr.ChatInterface(
    fn=chat_with_zeus,
    title="⚡ Zeus AI Assistant",
    description="مرحباً بك! أنا زيوس، مساعدك الذكي. كيف يمكنني مساعدتك اليوم؟",
    examples=["من أنت؟", "اشرح لي الذكاء الاصطناعي ببساطة", "كيف أتعلم البرمجة؟"],
    cache_examples=False,
    # لاحظ أننا لم نضف retry_btn هنا لتجنب خطأ الإصدارات
)

# 6. التشغيل
if __name__ == "__main__":
    demo.launch(share=True)