# -------------------------------------------------
#  اسم الملف: hf_chat.py
# -------------------------------------------------
import os
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

# ------------------- الإعدادات -------------------
# غير الاسم إلى أي نموذج تريده على Hugging Face
MODEL_NAME = "meta-llama/Llama-2-7b-chat-hf"          # مثال لنموذج مفتوح المصدر
DEVICE      = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# ------------------- تحميل النموذج ----------------
print(f"🔄 تحميل النموذج {MODEL_NAME} على {DEVICE} …")
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model     = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    torch_dtype=torch.float16 if DEVICE.type == "cuda" else torch.float32,
    low_cpu_mem_usage=True
).to(DEVICE)

print("✅ النموذج جاهز!\n")

# ------------------- حلقة الدردشة -----------------
def generate(prompt: str, max_new_tokens: int = 150, temperature: float = 0.7):
    inputs = tokenizer(prompt, return_tensors="pt").to(DEVICE)

    with torch.no_grad():
        output = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            do_sample=True,
            temperature=temperature,
            top_p=0.9,
            pad_token_id=tokenizer.eos_token_id,
        )
    # إرجاع النص المتولد (بدون الـ prompt الأصلي)
    text = tokenizer.decode(output[0], skip_special_tokens=True)
    return text[len(prompt):].strip()

if __name__ == "__main__":
    print("🗨️  اكتب جملة (أو اكتب 'exit' للخروج)")
    while True:
        user_input = input("\n>>> ")
        if user_input.lower() in {"exit", "quit"}:
            print("👋 وداعًا!")
            break

        answer = generate(user_input)
        print(f"\n🤖 {answer}")