import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from datasets import load_dataset
from peft import LoraConfig, get_peft_model
from trl import SFTTrainer
import os
import transformers

# --- 1. Параметрлерді орнату (CPU-ға Арналған) ---
MODEL_ID = "Qwen/Qwen1.5-0.5B-Chat" 
DATASET_FILE = "kz_book_data.jsonl" # Жаңа кітап дерегі
OUTPUT_DIR = "./qwen_kz_books_tuned"
NEW_TOKENIZER = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)

# --- 2. Деректерді Дайындау Функциясы ---
def formatting_prompts_func(examples):
    """Модельге арналған JSONL деректерін форматтау"""
    outputs = []
    for prompt, comp_kz, comp_mn in zip(examples['prompt'], examples['completion_kz'], examples['completion_mn']):
        
        # Екі тілді жауапты біріктіру (Монғолша бос болуы мүмкін)
        full_completion = f"Қазақша: {comp_kz}\n\nМонғолша: {comp_mn}"
        
        # Qwen-ге арналған чат үлгісін қолдану
        messages = [
            {"role": "user", "content": f"Сұрақ: {prompt}. Жауапты Қазақша етіп форматтаңыз:"},
            {"role": "assistant", "content": full_completion}
        ]
        
        # SFTTrainer үшін жаңа токенизатор арқылы форматтау
        input_text = NEW_TOKENIZER.apply_chat_template(
            messages, 
            tokenize=False
        )
        outputs.append(input_text)
    return {"text": outputs}

# --- 3. LoRA Конфигурациясы ---
lora_config = LoraConfig(
    r=8, 
    target_modules=["c_attn", "attn_q", "attn_k", "attn_v", "c_proj", "w1", "w2"],
    lora_alpha=8,
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM",
)

# --- 4. Модельді Жүктеу ---
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.float32, 
    device_map="cpu",
    trust_remote_code=True
)

# --- 5. Деректерді Жүктеу ---
try:
    dataset = load_dataset("json", data_files=DATASET_FILE, split="train")
except Exception as e:
    print(f"Деректер жиынтығын жүктеу қатесі: {e}. {DATASET_FILE} дұрыс JSONL форматында екеніне көз жеткізіңіз.")
    exit()

# --- 6. SFTTrainer-ді Қолдану ---
trainer = SFTTrainer(
    model=model,
    train_dataset=dataset,
    peft_config=lora_config,
    tokenizer=NEW_TOKENIZER,
    packing=True, 
    formatting_func=formatting_prompts_func,
    max_seq_length=512, 
    args=transformers.TrainingArguments(
        output_dir=OUTPUT_DIR,
        per_device_train_batch_size=1, 
        gradient_accumulation_steps=4,
        warmup_steps=5,
        max_steps=100, # Оқыту қадамдарын көбейттік
        learning_rate=2e-4,
        logging_steps=1,
        fp16=False,
        bf16=False, 
        optim="adamw_torch", 
        report_to="none",
    ),
)

# --- 7. Оқытуды Бастау ---
print("\n🔥 МОДЕЛЬДІ ОҚЫТУДЫ БАСТАЙМЫЗ (CPU-да жұмыс істейді)...")
trainer.train()

# --- 8. Нәтижені Сақтау ---
print("\n✅ Оқыту аяқталды. Нәтижелер сақталуда...")
trainer.model.save_pretrained(OUTPUT_DIR)
NEW_TOKENIZER.save_pretrained(OUTPUT_DIR)
print(f"Дайын модель {OUTPUT_DIR} папкасына сақталды.")