# %% Cell 2: Configuration & Auto-Detect GPU
import torch
import os
import shutil

# ============================
# PATHS (edit if needed)
# ============================
DATASET_PATH = r"E:\dataset\gemma4_ayurveda_unsloth_clean.jsonl"
OUTPUT_DIR   = r"E:\dataset\gemma4_ayurveda_a100_output"
MODEL_NAME   = "unsloth/gemma-4-12b-it"

# ============================
# GPU AUTO-DETECT
# ============================
if not torch.cuda.is_available():
    raise RuntimeError("No GPU found! Training requires CUDA.")

gpu_count = torch.cuda.device_count()
vram_gb = torch.cuda.get_device_properties(0).total_memory / 1024**3
print(f"GPUs visible : {gpu_count}")
print(f"GPU 0        : {torch.cuda.get_device_name(0)}")
print(f"VRAM         : {vram_gb:.1f} GB")

# Auto-configure based on actual VRAM
if vram_gb >= 35:
    # Full A100 40GB
    MAX_SEQ_LENGTH = 4096
    BATCH_SIZE = 2
    GRAD_ACCUM = 4
    LORA_R = 128
    LORA_ALPHA = 256
    LR = 2e-4
    print("\n✅ Detected: A100 40GB → Aggressive settings")
elif vram_gb >= 18:
    # MIG 20GB or similar
    MAX_SEQ_LENGTH = 2048
    BATCH_SIZE = 1
    GRAD_ACCUM = 8
    LORA_R = 64
    LORA_ALPHA = 128
    LR = 2e-4
    print("\n⚠️ Detected: ~20GB GPU → Conservative settings")
else:
    MAX_SEQ_LENGTH = 1024
    BATCH_SIZE = 1
    GRAD_ACCUM = 8
    LORA_R = 32
    LORA_ALPHA = 64
    LR = 3e-4
    print("\n⚠️ Detected: Low VRAM → Minimal settings")

EPOCHS = 1

# ============================
# VERIFY
# ============================
print(f"\n{'='*50}")
print("CONFIG")
print(f"{'='*50}")
print(f"Model         : {MODEL_NAME}")
print(f"Dataset       : {DATASET_PATH}")
print(f"Output        : {OUTPUT_DIR}")
print(f"Seq Length    : {MAX_SEQ_LENGTH}")
print(f"Batch         : {BATCH_SIZE} (eff: {BATCH_SIZE * GRAD_ACCUM})")
print(f"LoRA          : r={LORA_R}, alpha={LORA_ALPHA}")
print(f"LR            : {LR}")
print(f"Epochs        : {EPOCHS}")
print(f"{'='*50}")

if not os.path.exists(DATASET_PATH):
    print(f"\n❌ Dataset NOT found: {DATASET_PATH}")
else:
    ds_size = os.path.getsize(DATASET_PATH) / 1024**2
    print(f"\n✅ Dataset found: {ds_size:.1f} MB")

print("\n→ Run Cell 3 to load the dataset.")