"""
MasryGPT_chat Inference Example — English + مصري Bilingual
مثال تشغيل MasryGPT_chat — انجليزي ومصري
Model: ISLAM-PO/MasryGPT_chat (Qwen2.5-1.5B Egyptian Dialect)
Path: C:\Users\RDP\Downloads\MasryGPT-merged\MasryGPT-merged (1)\kaggle\working\MasryGPT-merged
"""

# English: Install requirements
# مصري: سطب المكتبات المطلوبة
# pip install -r requirements.txt  # or: pip install unsloth transformers accelerate bitsandbytes

# ========== Option 1: Unsloth (Fastest, Recommended) / أسرع طريقة ==========
from unsloth import FastLanguageModel

# English: Load merged 16-bit model in 4-bit for fast inference on T4
# مصري: حمّل الموديل المدمج 16-bit في وضع 4-bit عشان يشتغل بسرعة على T4
model, tokenizer = FastLanguageModel.from_pretrained(
    model_name = "ISLAM-PO/MasryGPT_chat",  # or local: "C:/.../MasryGPT-merged"
    # For local: model_name = r"C:\Users\RDP\Downloads\MasryGPT-merged\MasryGPT-merged (1)\kaggle\working\MasryGPT-merged"
    max_seq_length = 2048,
    dtype = None,
    load_in_4bit = True,  # False = 9-11GB VRAM, True = 2-4GB VRAM
)
FastLanguageModel.for_inference(model)  # 2x faster inference

# English: Chat template - system prompt forces Egyptian
# مصري: قالب الشات - رسالة النظام بتجبره يتكلم مصري
messages = [
    {"role": "system", "content": "انت مساعد مصري تتكلم عامية مصرية فقط، ردك قصير وطبيعي زي ولاد البلد"},
    {"role": "user", "content": "ازيك يا باشا؟ احكيلي نكتة مصرية قصيرة"},
    # Try also: "يعني ايه عيش وملح بين الصحاب؟" / "اشرح لخواجة يعني ايه يخرم بيتك"
]

# English: Apply chat template with attention mask to fix warning
# مصري: طبق قالب الشات مع attention_mask عشان تحذير Unsloth يختفي
inputs = tokenizer.apply_chat_template(
    messages,
    tokenize=True,
    add_generation_prompt=True,
    return_tensors="pt",
    return_dict=True,  # important to get attention_mask
).to("cuda")

# English: Generate - temperature 0.8 for creativity, 0.3 for facts
# مصري: التوليد - 0.8 للهزار والإبداع، 0.3 للمعلومات الجادة
outputs = model.generate(
    **inputs,
    max_new_tokens=256,      # English: increase for longer jokes / مصري: زوده للنكت الطويلة
    temperature=0.8,
    top_p=0.95,
    do_sample=True,
    repetition_penalty=1.1,  # reduces repetition / يقلل التكرار
    use_cache=True,
)

response = tokenizer.batch_decode(outputs[:, inputs["input_ids"].shape[1]:], skip_special_tokens=True)[0]
print("="*40)
print("MasryGPT Response / رد MasryGPT:")
print(response)
print("="*40)


# ========== Option 2: Transformers (Standard) / الطريقة العادية ==========
"""
from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained(
    "ISLAM-PO/MasryGPT_chat",
    device_map="auto",
    trust_remote_code=True,
)
tokenizer = AutoTokenizer.from_pretrained("ISLAM-PO/MasryGPT_chat", fix_mistral_regex=True)
# same generate code as above
"""

# ========== Quick Test Prompts / أسئلة سريعة للتجربة ==========
test_prompts = [
    "ازيك عامل ايه؟",
    "احكيلي نكتة مصرية",
    "يعني ايه عيش وملح؟",
    "اشرح يعني ايه 'على راسي' بالمصري",
    "اعملي بوست فيسبوك مصري عن رمضان",
]
# Loop them with same messages logic