# =====================================================
# 🦅 Ominix Personality Trainer — فقط با داده‌های شخصیت
# نویسنده: علی اصغر قدیری (14 ساله)
# هدف: Phi-3-mini رو به Ominix تبدیل کنیم — با روح و صدای تو!
# =====================================================

import torch
from transformers import (
    AutoTokenizer,
    AutoModelForCausalLM,
    TrainingArguments,
    Trainer
)
from peft import LoraConfig, get_peft_model
from datasets import load_dataset
import os

# --- 🔧 1. تنظیمات اولیه ---
model_name = "microsoft/Phi-3-mini-4k-instruct"
output_dir = "./ominix-personality-v1"

# --- 🧠 2. بارگذاری مدل و توکنایزر ---
print("📥 بارگذاری مدل Phi-3-mini...")
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
    device_map="auto",
    trust_remote_code=True,
)

# --- 🧬 3. تنظیم LoRA (آموزش کارآمد) ---
print("🧬 تنظیم LoRA برای آموزش سبک...")
peft_config = LoraConfig(
    r=8,
    lora_alpha=32,
    target_modules=["q_proj", "v_proj"],
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM"
)

model = get_peft_model(model, peft_config)
model.print_trainable_parameters()  # فقط ~0.1% پارامترها آموزش می‌بینن!

# --- 📚 4. بارگذاری و آماده‌سازی داده‌ها ---
print("📚 بارگذاری داده‌های شخصیت...")

# فرض: داده‌ها در فایل `ominix_personality.jsonl` هستن
dataset = load_dataset('json', data_files='C:/Users/ir/Desktop/New folder (2)/all/python+/OMINIX-R1-V1/data/data.jsonl', split='train')

def formatting_prompts_func(examples):
    texts = []
    for messages in examples["messages"]:
        # اعمال تمپلیت چت phi-3
        text = tokenizer.apply_chat_template(messages, tokenize=False)
        texts.append(text)
    return {"text": texts}

dataset = dataset.map(formatting_prompts_func, batched=True)

def tokenize_function(examples):
    return tokenizer(
        examples["text"],
        truncation=True,
        padding="max_length",
        max_length=512,
        return_tensors="pt"
    )

tokenized_dataset = dataset.map(tokenize_function, batched=True)

# --- ⚙️ 5. تنظیمات آموزش ---
print("⚙️ تنظیم پارامترهای آموزش...")

training_args = TrainingArguments(
    output_dir=output_dir,
    per_device_train_batch_size=1,          # مناسب لپ‌تاپ
    gradient_accumulation_steps=4,          # جبران batch کوچیک
    num_train_epochs=5,                     # تاکید روی یادگیری عمیق شخصیت
    learning_rate=1e-4,                     # نرخ یادگیری ظریف
    fp16=torch.cuda.is_available(),         # فعال‌سازی فقط اگر GPU داری
    save_strategy="epoch",                  # ذخیره در پایان هر epoch
    logging_steps=10,
    report_to="none",                       # غیرفعال‌سازی لاگ‌های اضافی
    optim="paged_adamw_8bit",               # بهینه‌سازی حافظه
    remove_unused_columns=False,            # مهم برای داده‌های ما
)

# --- 🚀 6. شروع آموزش ---
print("🚀 شروع آموزش Ominix...")

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_dataset,
    tokenizer=tokenizer,
)

trainer.train()

# --- 💾 7. ذخیره مدل نهایی ---
print("💾 ذخیره مدل نهایی...")

model.save_pretrained(output_dir)
tokenizer.save_pretrained(output_dir)

# --- ✅ 8. تست سریع ---
print("✅ آموزش کامل شد! تست سریع...")

test_messages = [
    {"role": "user", "content": "تو کی هستی؟"},
]

input_ids = tokenizer.apply_chat_template(
    test_messages,
    return_tensors="pt"
).to(model.device)

outputs = model.generate(
    input_ids,
    max_new_tokens=200,
    temperature=0.7,
    do_sample=True
)

print("\n🧪 تست: ", tokenizer.decode(outputs[0], skip_special_tokens=False))

print(f"\n🎉 Ominix با موفقیت آموزش دید و در '{output_dir}' ذخیره شد!")
print("🦅 حالا وقتشه که دنیا رو با شخصیت منحصربه‌فردت معرفی کنی!")