# ==========================================
# CELL 1: CONFIGURATION
# ==========================================
# Pick your model. For A4000 16GB use 5B. For 24GB+ use 12B.
MODEL_NAME = "unsloth/gemma-4-E2B-it"      # ~5B params, fits 16GB
# MODEL_NAME = "unsloth/gemma-4-12b-it"    # ~12B params, needs 24GB+
# MODEL_NAME = "unsloth/gemma-4-27b-it"    # ~27B params, needs 40GB+

# Dataset path (use 1.69GB clean file for full training, or 404MB for quick tests)
DATASET_PATH = r"E:\dataset\gemma4_ayurveda_unsloth_clean.jsonl"

# Where to save the trained LoRA adapter
OUTPUT_DIR = r"E:\dataset\gemma4_ayurveda_lora_output"

# Training settings
MAX_SEQ_LENGTH = 2048        # 2048 for 16GB safe. Use 3072/4096 if you have more VRAM.
BATCH_SIZE = 1               # Keep at 1 for 16GB. Increase to 2 if you have 24GB+.
GRAD_ACCUM = 8               # Effective batch = BATCH_SIZE * GRAD_ACCUM
EPOCHS = 1                   # 1 epoch is enough for 971k conversations
LORA_R = 64                  # Rank: 64 for deep domain (Sanskrit/medical). 32 for quick tests.
LORA_ALPHA = 128             # Always 2x rank
LR = 3e-4                    # Learning rate: 3e-4 for 5B, 2e-4 for 12B/27B

import os
print("=" * 60)
print("CONFIGURATION")
print("=" * 60)
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} (effective: {BATCH_SIZE * GRAD_ACCUM})")
print(f"LoRA       : r={LORA_R}, alpha={LORA_ALPHA}")
print(f"Epochs     : {EPOCHS}")
print(f"LR         : {LR}")
print("=" * 60)
print(f"Dataset exists: {os.path.exists(DATASET_PATH)}")