app.py

# -*- coding: utf-8 -*-

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

model_path = "inceptionai/jais-family-13b-chat"

# Jais original prompts (مثل الكود الأصلي)
prompt_eng = "### Instruction:Your name is 'Jais', and you are named after Jebel Jais, the highest mountain in UAE. You were made by 'Inception' in the UAE. You are a helpful, respectful, and honest assistant. Always answer as helpfully as possible, while being safe. Complete the conversation between [|Human|] and [|AI|]:\n### Input: [|Human|] {Question}\n[|AI|]\n### Response :"
prompt_ar = "### Instruction:اسمك \"جيس\" وسميت على اسم جبل جيس اعلى جبل في الامارات. تم بنائك بواسطة Inception في الإمارات. أنت مساعد مفيد ومحترم وصادق. أجب دائمًا بأكبر قدر ممكن من المساعدة، مع الحفاظ على البقاء أمناً. أكمل المحادثة بين [|Human|] و[|AI|] :\n### Input:[|Human|] {Question}\n[|AI|]\n### Response :"

# إذا كان فيه HF_TOKEN في البيئة
hf_token = os.getenv("HF_TOKEN")

device = "cuda" if torch.cuda.is_available() else "cpu"

# تحميل المودل مع تحسينات لـ ZeroGPU
print("جاري تحميل المودل...")
tokenizer = AutoTokenizer.from_pretrained(
    model_path,
    token=hf_token,
    trust_remote_code=True
)

model = AutoModelForCausalLM.from_pretrained(
    model_path,
    device_map=None,  # إزالة device_map لتجنب مشاكل مع past_key_values
    trust_remote_code=True,
    token=hf_token,
    torch_dtype=torch.float16 if device == "cuda" else torch.float32,
    low_cpu_mem_usage=True
).to(device)  # نقل المودل إلى الجهاز المحدد
print("تم تحميل المودل بنجاح!")

if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

def get_response(text, tokenizer=tokenizer, model=model):
    """نفس الدالة من الكود الأصلي مع تحسينات للأداء"""
    input_ids = tokenizer(text, return_tensors="pt").input_ids
    inputs = input_ids.to(device)
    input_len = inputs.shape[-1]
    generate_ids = model.generate(
        inputs,
        top_p=0.9,
        temperature=0.3,
        max_length=2048,
        min_length=input_len + 4,
        repetition_penalty=1.2,
        do_sample=True,
        pad_token_id=tokenizer.pad_token_id if tokenizer.pad_token_id is not None else tokenizer.eos_token_id
    )
    response = tokenizer.batch_decode(
        generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True
    )[0]
    response = response.split("### Response :")[-1]
    return response

@spaces.GPU()
def generate_response(input_data, chat_history):
    # Use Arabic prompt for all responses (simplified approach)
    ques = input_data
    text = prompt_ar.format_map({'Question': ques})

    try:
        response = get_response(text)
        yield response

    except Exception as e:
        print(f"Error in generate_response: {e}")
        import traceback
        print(traceback.format_exc())
        yield "أعتذر، حدث خطأ. يرجى المحاولة مرة أخرى."

demo = gr.ChatInterface(
    fn=generate_response,
    examples=[
        [{"text": "ما هي عاصمة الامارات؟"}],
        [{"text": "شرح لي الذكاء الاصطناعي"}],
        [{"text": "أخبرني قصة قصيرة"}],
        [{"text": "كيف أتعلم البرمجة؟"}],
        [{"text": "ما هو الفرق بين الذكاء الاصطناعي والتعلم الآلي؟"}],
    ],
    cache_examples=False,
    type="messages",
    title="Jais AI - ذكاء جيس العربي",
    description="""🤖 ذكاء صناعي متقدم يدعم اللغة العربية

✨ المميزات:
- 🇸🇦 التركيز على اللغة العربية
- 🧠 موديل Jais المتقدم من Inception
- 💬 إجابات ذكية واحترافية
- 🎯 سهل الاستخدام بدون إعدادات معقدة

احجي مع ذكاء Jais الاصطناعي بالعربية.""",
    fill_height=True,
    textbox=gr.Textbox(
        label="اكتب رسالتك هنا",
        placeholder="مثال: شرح لي موضوع معقد..."
    ),
    stop_btn="إيقاف التوليد",
    multimodal=False,
    theme=gr.themes.Soft()
)

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