import os
from unsloth import FastLanguageModel
from datasets import load_dataset
from trl import SFTTrainer
from transformers import TrainingArguments

# 1. AYARLAR
max_seq_length = 2048 # Modelin tek seferde okuyacağı maksimum kelime/token sayısı
dtype = None # Ekran kartına göre otomatik ayarlar (Float16 / Bfloat16)
load_in_4bit = True # VRAM dostu 4-bit sıkıştırma aktif

# 2. MODELİ YÜKLE
# Buraya Hugging Face'e yüklediğin kendi modelinin yolunu yazabilirsin
model, tokenizer = FastLanguageModel.from_pretrained(
    model_name = "tda45/TdAI-4bit", 
    max_seq_length = max_seq_length,
    dtype = dtype,
    load_in_4bit = True, # Bu değerin True olduğundan kesinlikle emin ol
    fix_tokenizer = True,
)

# 3. LoRA AYARLARI (Eğitimi hızlandıran akıllı katmanlar)
model = FastLanguageModel.get_peft_model(
    model,
    r = 16, # LoRA matris boyutu (Önerilen: 8, 16, 32)
    target_modules = ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
    lora_alpha = 16,
    lora_dropout = 0, 
    bias = "none",    
    use_gradient_checkpointing = "unsloth", # Hafıza tasarrufu için önemli
    random_state = 3407,
    use_rslora = False,
    loftq_config = None,
)

# 4. VERİ SETİNİ HAZIRLA
# Alpaca formatında örnek bir prompt şablonu
prompt_sablonu = """Aşağıda bir görevi tanımlayan bir talimat ve daha fazla ortam sağlayan bir girdi bulunmaktadır. İsteyi uygun şekilde tamamlayan bir yanıt yazın.

### Talimat:
{}

### Yanıt:
{}"""

# Eğer Hugging Face'te veya yerelde bir veri setin varsa buraya bağla
# Örnek olarak kendi json dosyanı yüklemek istersen: dataset = load_dataset("json", data_files="verilerim.json", split="train")
dataset = load_dataset("yahma/alpaca-cleaned", split = "train") # Örnek hazır veri seti

def format_prompts(examples):
    instructions = examples["instruction"]
    outputs      = examples["output"]
    texts = []
    for instruction, output in zip(instructions, outputs):
        text = prompt_sablonu.format(instruction, output) + tokenizer.eos_token
        texts.append(text)
    return { "text" : texts, }

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

# 5. EĞİTİM PARAMETRELERİ
trainer = SFTTrainer(
    model = model,
    tokenizer = tokenizer,
    train_dataset = dataset,
    dataset_text_field = "text",
    max_seq_length = max_seq_length,
    dataset_num_proc = 2,
    packing = False, # Küçük veri setleri için False kalması daha iyi
    args = TrainingArguments(
        per_device_train_batch_size = 2,
        gradient_accumulation_steps = 4,
        warmup_steps = 5,
        max_steps = 60, # Deneme için adım sayısını küçük tuttuk, duruma göre artırabilirsin
        learning_rate = 2e-4,
        fp16 = not FastLanguageModel.is_bfloat16_supported(),
        bf16 = FastLanguageModel.is_bfloat16_supported(),
        logging_steps = 1,
        optim = "adamw_8bit",
        weight_decay = 0.01,
        lr_scheduler_type = "linear",
        seed = 3407,
        output_dir = "outputs",
    ),
)

# 6. EĞİTİMİ BAŞLAT
print("Eğitim süreci başlıyor...")
trainer_stats = trainer.train()
print("Eğitim tamamlandı!")

# 7. EĞİTİLMİŞ MODELİ KAYDET
model.save_pretrained("lora_model")
tokenizer.save_pretrained("lora_model")
print("Eğitilen LoRA katmanları 'lora_model' klasörüne kaydedildi.")