#!/usr/bin/env python3
# ============================================================
# ZabaanAI-v2: CPU-Only Training Script (Gradient Checkpointing + LoRA)
# WARNING: This will be VERY SLOW on CPU (~10-100x slower than GPU)
# For 7B models on CPU: only for small experiments or 1B models
# Recommended: use google/gemma-2b or microsoft/phi-2 for CPU training
# ============================================================
import os, warnings, sys
warnings.filterwarnings('ignore')

from transformers import (
    AutoTokenizer, AutoModelForCausalLM,
    TrainingArguments, set_seed
)
from peft import LoraConfig, get_peft_model, TaskType, prepare_model_for_kbit_training
from datasets import load_dataset
import torch

MODEL_ID   = 'microsoft/phi-2'  # 2.7B = smallest model that trains on CPU RAM
HF_REPO    = 'shaikhsalman/zabaanai-v2-cpu'
DATA_PATH  = '/app/data/processed/zabaanai_v2_sft_pakistan.jsonl'
OUTPUT_DIR = '/app/output/zabaanai-v2-cpu'

LORA_RANK  = 16
SEQ_LEN    = 512          # reduced for CPU RAM
BS         = 1
GRAD_ACCUM = 32            # effective batch = 32
EPOCHS     = 2
LR         = 1e-3          # LoRA default lr

if not os.path.exists(DATA_PATH):
    print(f'Data not found: {DATA_PATH} — run 01_curate_sft_data.py first!')
    sys.exit(1)

set_seed(42)
print('Loading tokenizer...')
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
tokenizer.pad_token = tokenizer.eos_token

print('Loading dataset...')
ds = load_dataset('json', data_files=DATA_PATH, split='train')
ds = ds.train_test_split(test_size=0.02, seed=42)
print(f'  Train: {len(ds[\"train\"])} | Eval: {len(ds[\"test\"])}')

def preprocess(example):
    text = example.get('text', '')
    enc = tokenizer(text, max_length=SEQ_LEN, truncation=True,
                    padding='max_length', return_tensors=None)
    enc['labels'] = enc['input_ids'].copy()
    return enc

print('Tokenizing...')
tokenized = ds.map(preprocess, remove_columns=ds.column_names, num_proc=2)

print('Loading model with LoRA...')
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    device_map='cpu',
    torch_dtype=torch.float16,
    offload_folder='/tmp/offload',
    trust_remote_code=True,
)
model = prepare_model_for_kbit_training(model)
model.config.use_cache = False

lora_config = LoraConfig(
    r=LORA_RANK, lora_alpha=LORA_RANK*2,
    lora_dropout=0.05,
    target_modules=['q_proj','k_proj','v_proj','o_proj','gate_proj','up_proj','down_proj'],
    bias='none', task_type=TaskType.CAUSAL_LM,
)
model = get_peft_model(model, lora_config)
model.print_trainable_parameters()

from trl import SFTTrainer
trainer = SFTTrainer(
    model=model,
    args=TrainingArguments(
        output_dir=OUTPUT_DIR,
        per_device_train_batch_size=BS,
        per_device_eval_batch_size=BS,
        gradient_accumulation_steps=GRAD_ACCUM,
        num_train_epochs=EPOCHS,
        learning_rate=LR,
        bf16=False, fp16=True,
        logging_dir=f'{OUTPUT_DIR}/logs',
        logging_strategy='steps', logging_steps=10,
        save_strategy='steps', save_steps=200,
        eval_steps=200, eval_strategy='steps',
        save_total_limit=2,
        report_to=['none'],
        dataloader_num_workers=1,
        seed=42,
        max_grad_norm=1.0,
        gradient_checkpointing=True,
        gradient_checkpointing_kwargs={'use_reentrant': False},
    ),
    train_dataset=tokenized['train'],
    eval_dataset=tokenized['test'],
    data_collator=None,
    max_seq_length=SEQ_LEN,
    tokenizer=tokenizer,
    dataset_text_field='text',
)

print('Starting CPU training (WARNING: this will take many hours/days)...')
print(f'  Effective batch: {BS * GRAD_ACCUM}, steps: ~{len(tokenized[\"train\"]) // GRAD_ACCUM * EPOCHS}')
trainer.train()
trainer.save_model(f'{OUTPUT_DIR}/final')
print('CPU training complete!')