#!/usr/bin/env python3
# %% Cell 1 - Setup & Imports
"""
Jupyter Notebook Script for Training Gemma-4-E2B-it on Ayurveda Dataset
Optimized for: RTX A4000 16GB VRAM + 20GB System RAM
Uses: Unsloth (pip version), 1.69GB clean dataset
"""

import os
import gc
import torch

# ============================================================================
# CONFIGURATION - EDIT THESE PATHS IF NEEDED
# ============================================================================
# The 1.69GB clean dataset (messages only, no metadata, role=assistant)
DATASET_PATH = r"E:\dataset\gemma4_ayurveda_unsloth_clean.jsonl"

# Model name - will load from local HuggingFace cache
# Local cache path: C:\Users\Antplay\.cache\huggingface\hub\models--unsloth--gemma-4-E2B-it
MODEL_NAME = "unsloth/gemma-4-E2B-it"

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

# Training hyperparameters (tuned for A4000 16GB)
MAX_SEQ_LENGTH = 2048        # Use 2048 to stay safely under 16GB VRAM. Change to 3072 if you have headroom.
BATCH_SIZE = 1               # MUST be 1 for 16GB with context 2048+
GRAD_ACCUM = 8               # Effective batch = 8
LORA_R = 64                  # High rank for domain knowledge
LORA_ALPHA = 128             # 2x rank
LEARNING_RATE = 3e-4         # Slightly aggressive for 5B model
NUM_EPOCHS = 1
WARMUP_STEPS = 500

print("=" * 60)
print("CONFIGURATION")
print("=" * 60)
print(f"Dataset    : {DATASET_PATH}")
print(f"Model      : {MODEL_NAME}")
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     : {NUM_EPOCHS}")
print("=" * 60)

# Verify dataset exists
if not os.path.exists(DATASET_PATH):
    print(f"ERROR: Dataset not found at {DATASET_PATH}")
    print("Please check the path and update DATASET_PATH above.")
    raise FileNotFoundError(f"Dataset not found: {DATASET_PATH}")
else:
    size_mb = os.path.getsize(DATASET_PATH) / (1024 * 1024)
    print(f"Dataset found: {size_mb:.1f} MB")

# Verify CUDA
if not torch.cuda.is_available():
    print("WARNING: CUDA not available. Training will be extremely slow on CPU.")
else:
    print(f"CUDA available: {torch.cuda.get_device_name(0)}")
    print(f"VRAM total    : {torch.cuda.get_device_properties(0).total_memory / 1024**3:.1f} GB")
    print(f"VRAM current  : {torch.cuda.memory_allocated() / 1024**3:.1f} GB")

print("\nCell 1 complete. Run Cell 2 to load the dataset.")


# %% Cell 2 - Load Dataset (memory efficient)
print("\n" + "=" * 60)
print("[Cell 2] Loading Dataset...")
print("=" * 60)

from datasets import load_dataset

# Load the JSONL dataset
# The clean file is 1.69GB -> ~2-3GB in RAM after loading
dataset = load_dataset(
    "json",
    data_files=DATASET_PATH,
    split="train",
    streaming=False,  # Standardize_sharegpt needs a regular Dataset
)

print(f"Dataset loaded: {len(dataset):,} conversations")

# Clear memory before standardization
gc.collect()

print("\nCell 2 complete. Run Cell 3 to standardize format.")


# %% Cell 3 - Standardize ShareGPT Format
print("\n" + "=" * 60)
print("[Cell 3] Standardizing ShareGPT format...")
print("=" * 60)

from unsloth import standardize_sharegpt

# This converts the messages array into a text field that the model can train on
# It adds a "text" column to the dataset
dataset = standardize_sharegpt(dataset)

print(f"Standardization complete.")
print(f"Sample text length: {len(dataset[0]['text'])} characters")
print(f"Sample preview:\n{dataset[0]['text'][:300]}...")

# Free memory
gc.collect()

print("\nCell 3 complete. Run Cell 4 to load the model.")


# %% Cell 4 - Load Model from Local Cache
print("\n" + "=" * 60)
print("[Cell 4] Loading Gemma-4-E2B-it from local cache...")
print("=" * 60)

from unsloth import FastLanguageModel

# local_files_only=True forces loading from cache (no download)
model, tokenizer = FastLanguageModel.from_pretrained(
    model_name=MODEL_NAME,
    max_seq_length=MAX_SEQ_LENGTH,
    dtype=None,              # Auto-detect: bf16 on Ampere (A4000)
    load_in_4bit=True,       # QLoRA 4-bit
    local_files_only=True,   # Use your already-downloaded model
)

print(f"Model loaded successfully.")
print(f"VRAM after model load: {torch.cuda.memory_allocated() / 1024**3:.1f} GB")

gc.collect()
torch.cuda.empty_cache()

print("\nCell 4 complete. Run Cell 5 to attach LoRA adapter.")


# %% Cell 5 - Attach LoRA Adapter
print("\n" + "=" * 60)
print("[Cell 5] Attaching QLoRA adapter...")
print("=" * 60)

model = FastLanguageModel.get_peft_model(
    model,
    r=LORA_R,
    lora_alpha=LORA_ALPHA,
    lora_dropout=0,        # 0: 971k samples = no overfitting risk
    bias="none",
    use_rslora=False,
    use_gradient_checkpointing="unsloth",  # Saves ~30% VRAM
    random_state=3407,
    target_modules=[
        "q_proj", "k_proj", "v_proj", "o_proj",
        "gate_proj", "up_proj", "down_proj",
    ],
)

print(f"LoRA adapter attached.")
print(f"Trainable parameters: {model.print_trainable_parameters()}")
print(f"VRAM after adapter: {torch.cuda.memory_allocated() / 1024**3:.1f} GB")

gc.collect()
torch.cuda.empty_cache()

print("\nCell 5 complete. Run Cell 6 to configure training.")


# %% Cell 6 - Training Arguments
print("\n" + "=" * 60)
print("[Cell 6] Configuring training arguments...")
print("=" * 60)

from transformers import TrainingArguments

# Calculate total steps for reference
total_steps = len(dataset) // (BATCH_SIZE * GRAD_ACCUM) * NUM_EPOCHS
print(f"Estimated total steps: {total_steps:,}")
print(f"Estimated time on A4000: ~{total_steps / 2.5 / 3600:.1f} hours (at ~2.5 it/s)")

training_args = TrainingArguments(
    output_dir=OUTPUT_DIR,
    overwrite_output_dir=True,

    num_train_epochs=NUM_EPOCHS,
    per_device_train_batch_size=BATCH_SIZE,
    gradient_accumulation_steps=GRAD_ACCUM,

    learning_rate=LEARNING_RATE,
    lr_scheduler_type="cosine",
    warmup_steps=WARMUP_STEPS,

    optim="adamw_8bit",
    weight_decay=0.01,
    max_grad_norm=0.3,

    bf16=torch.cuda.is_bf16_supported(),  # True on A4000
    fp16=not torch.cuda.is_bf16_supported(),

    logging_steps=50,
    save_strategy="steps",
    save_steps=2000,
    save_total_limit=2,

    group_by_length=True,
    remove_unused_columns=False,

    report_to="none",  # Change to "wandb" if you want tracking
    seed=3407,
)

print("Training arguments configured.")
print(f"  Effective batch size: {BATCH_SIZE * GRAD_ACCUM}")
print(f"  Learning rate: {LEARNING_RATE}")
print(f"  Context length: {MAX_SEQ_LENGTH}")
print(f"  Epochs: {NUM_EPOCHS}")

print("\nCell 6 complete. Run Cell 7 to initialize the trainer.")


# %% Cell 7 - Initialize SFT Trainer
print("\n" + "=" * 60)
print("[Cell 7] Initializing SFTTrainer...")
print("=" * 60)

from trl import SFTTrainer

trainer = SFTTrainer(
    model=model,
    tokenizer=tokenizer,
    train_dataset=dataset,
    dataset_text_field="text",
    max_seq_length=MAX_SEQ_LENGTH,
    args=training_args,
    packing=False,  # NEVER pack instruction data (system/user/assistant boundaries break)
)

print("Trainer initialized successfully.")
print(f"VRAM before training: {torch.cuda.memory_allocated() / 1024**3:.1f} GB")

gc.collect()
torch.cuda.empty_cache()

print("\nCell 7 complete. Run Cell 8 to START TRAINING.")


# %% Cell 8 - Train!
print("\n" + "=" * 60)
print("[Cell 8] STARTING TRAINING")
print("=" * 60)
print(f"Dataset conversations : {len(dataset):,}")
print(f"Epochs              : {NUM_EPOCHS}")
print(f"Batch size          : {BATCH_SIZE} (effective: {BATCH_SIZE * GRAD_ACCUM})")
print(f"Context length      : {MAX_SEQ_LENGTH}")
print(f"Learning rate       : {LEARNING_RATE} (cosine)")
print(f"LoRA rank           : {LORA_R}")
print(f"Expected time       : ~{total_steps / 2.5 / 3600:.1f} hours")
print("=" * 60)

trainer.train()

print("\nTraining complete!")
print(f"Final loss: {trainer.state.log_history[-1].get('loss', 'N/A')}")

print("\nCell 8 complete. Run Cell 9 to save.")


# %% Cell 9 - Save LoRA Adapter
print("\n" + "=" * 60)
print("[Cell 9] Saving trained adapter...")
print("=" * 60)

import os

os.makedirs(OUTPUT_DIR, exist_ok=True)

# Save adapter only (small, ~100-150 MB)
adapter_path = os.path.join(OUTPUT_DIR, "adapter")
model.save_pretrained(adapter_path)
tokenizer.save_pretrained(adapter_path)

print(f"Adapter saved to: {adapter_path}")
print(f"Size: {sum(os.path.getsize(os.path.join(adapter_path, f)) for f in os.listdir(adapter_path)) / 1024**2:.1f} MB")

# Optional: Merge adapter into base model (requires ~12GB VRAM temporarily)
print("\nAttempting optional merge into 16-bit model...")
try:
    merged_path = os.path.join(OUTPUT_DIR, "merged_16bit")
    model.save_pretrained_merged(merged_path, tokenizer, save_method="merged_16bit")
    print(f"Merged model saved to: {merged_path}")
except RuntimeError as e:
    print(f"Merge skipped (likely OOM during merge): {e}")
    print("You can still use the adapter + base model for inference.")

gc.collect()
torch.cuda.empty_cache()

print("\n" + "=" * 60)
print("ALL DONE!")
print("=" * 60)
print(f"Output folder: {OUTPUT_DIR}")
print(f"To use your trained model, load the adapter from: {adapter_path}")
print("\nHappy training!")


# %% Cell 10 - Quick Inference Test (Optional)
print("\n" + "=" * 60)
print("[Cell 10] Quick Inference Test")
print("=" * 60)

test_messages = [
    {"role": "user", "content": "What does Charaka say about the qualities of a good physician in Sutrasthana?"}
]

inputs = tokenizer.apply_chat_template(test_messages, tokenize=True, return_tensors="pt", add_generation_prompt=True).to("cuda")

print("Generating response...")
outputs = model.generate(
    inputs,
    max_new_tokens=512,
    temperature=0.7,
    top_p=0.9,
    do_sample=True,
)

response = tokenizer.decode(outputs[0], skip_special_tokens=True)
print("\nResponse:")
print(response)