#!/usr/bin/env python3
"""
اختبار الموديل العراقي 4-bit
"""

import os
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

def test_model():
    model_path = "anaspro/iraqi-7b"
    hf_token = os.getenv("HF_TOKEN")

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

    model = AutoModelForCausalLM.from_pretrained(
        model_path,
        device_map="auto",
        trust_remote_code=True,
        token=hf_token,
        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
        low_cpu_mem_usage=True
    )

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

    # اختبار الـ prompt الجديد
    prompt = """### Instruction:أنت ذكاء صناعي يتحدث باللهجة العراقية ويجيب باحترافية وذكاء. الإجابات دائما تستخدم اللهجة العراقية.
### Input:[|Human|] مرحبا كيفك؟
[|AI|]
### Response :"""

    inputs = tokenizer(prompt, return_tensors="pt")
    input_ids = inputs['input_ids'].to('cuda' if torch.cuda.is_available() else 'cpu')
    attention_mask = inputs['attention_mask'].to('cuda' if torch.cuda.is_available() else 'cpu')

    print("جاري توليد الرد...")
    outputs = model.generate(
        input_ids,
        attention_mask=attention_mask,
        max_length=input_ids.shape[-1] + 100,
        temperature=0.3,
        top_p=0.8,
        do_sample=True,
        pad_token_id=tokenizer.pad_token_id,
        eos_token_id=tokenizer.eos_token_id
    )

    response = tokenizer.decode(outputs[0], skip_special_tokens=True)

    # استخراج الرد الجديد فقط
    if "### Response :" in response:
        response = response.split("### Response :")[-1].strip()

    print("=" * 50)
    print("الاختبار نجح! ✅")
    print(f"الموديل: {model_path}")
    print("الـ prompt: ذكاء صناعي باللهجة العراقية (مع التأكيد على استخدام اللهجة دائماً)"
    print(f"الرد: {response}")
    print("=" * 50)

if __name__ == "__main__":
    test_model()