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

# 1. تحديد النموذج (استخدمنا 2b لأنه الأنسب للمساحات المجانية وسريع)
model_name = "google/gemma-2-2b-it"

# 2. الحصول على التوكن من إعدادات المساحة (Secrets)
token = os.getenv("HF_TOKEN")

# 3. تحميل المحلل اللفظي والنموذج
tokenizer = AutoTokenizer.from_pretrained(model_name, token=token)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.bfloat16, # أفضل للأداء والدقة
    device_map="auto",
    token=token
)

def chat_function(message, history):
    # بناء سجل المحادثة بتنسيق Gemma
    messages = []
    for user_msg, assistant_msg in history:
        messages.append({"role": "user", "content": user_msg})
        messages.append({"role": "assistant", "content": assistant_msg})
    
    messages.append({"role": "user", "content": message})

    # تجهيز المدخلات باستخدام Template النموذج الرسمي
    input_ids = tokenizer.apply_chat_template(
        messages,
        add_generation_prompt=True,
        return_tensors="pt"
    ).to(model.device)

    # توليد الرد
    outputs = model.generate(
        input_ids,
        max_new_tokens=512,
        do_sample=True,
        temperature=0.7,
        top_p=0.9,
    )
    
    # فك التشفير واستخراج النص الجديد فقط
    response = tokenizer.decode(outputs[0][input_ids.shape[-1]:], skip_special_tokens=True)
    return response

# 4. إنشاء واجهة الدردشة
demo = gr.ChatInterface(
    fn=chat_function,
    title="Gemma 2 Chatbot",
    description="دردشة مباشرة مع نموذج Gemma من جوجل على Hugging Face Spaces",
    examples=["كيف حالك؟", "اشرح لي الثقوب السوداء ببساطة", "اكتب قصيدة قصيرة عن الذكاء الاصطناعي"],
    cache_examples=False,
)

if __name__ == "__main__":
    demo.launch()