# ============================================================
# AYURVEDA GEMMA-4 TRAINING FOR A100
# All 10 Jupyter cells in one file
# Copy each section between # %% markers into a Jupyter cell
# ============================================================

# ============================================================
# CELL 1: Install Packages (Run once, restart kernel after)
# ============================================================
!pip install -q unsloth transformers datasets trl accelerate bitsandbytes
print("Installation complete. Restart the kernel, then run Cell 2.")


# ============================================================
# CELL 2: Configuration & Auto-Detect GPU
# ============================================================
import torch
import os
import shutil

# --- EDIT THESE PATHS 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")

if vram_gb >= 35:
    MAX_SEQ_LENGTH = 4096
    BATCH_SIZE = 2
    GRAD_ACCUM = 4
    LORA_R = 128
    LORA_ALPHA = 256
    LR = 2e-4
    print("\nDetected: A100 40GB -> Aggressive settings")
elif vram_gb >= 18:
    MAX_SEQ_LENGTH = 2048
    BATCH_SIZE = 1
    GRAD_ACCUM = 8
    LORA_R = 64
    LORA_ALPHA = 128
    LR = 2e-4
    print("\nDetected: ~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("\nDetected: Low VRAM -> Minimal settings")

EPOCHS = 1

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"\nWARNING: Dataset NOT found: {DATASET_PATH}")
else:
    ds_size = os.path.getsize(DATASET_PATH) / 1024**2
    print(f"\nDataset found: {ds_size:.1f} MB")

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


# ============================================================
# CELL 3: Load Dataset
# ============================================================
from datasets import load_dataset
from unsloth import standardize_sharegpt
import gc

dataset = load_dataset("json", data_files=DATASET_PATH, split="train")
print(f"Loaded {len(dataset):,} conversations")

dataset = standardize_sharegpt(dataset)
print(f"Standardized. Sample text: {len(dataset[0]['text'])} chars")

gc.collect()
print("\n-> Run Cell 4 to load the model.")


# ============================================================
# CELL 4: Load Gemma-4 12B Model
# ============================================================
from unsloth import FastLanguageModel
import gc

model, tokenizer = FastLanguageModel.from_pretrained(
    model_name=MODEL_NAME,
    max_seq_length=MAX_SEQ_LENGTH,
    dtype=None,
    load_in_4bit=True,
)

print(f"Model loaded: {MODEL_NAME}")
print(f"VRAM used: {torch.cuda.memory_allocated() / 1024**3:.1f} GB")

gc.collect()
torch.cuda.empty_cache()
print("\n-> Run Cell 5 to attach LoRA adapter.")


# ============================================================
# CELL 5: Attach LoRA Adapter
# ============================================================
model = FastLanguageModel.get_peft_model(
    model,
    r=LORA_R,
    lora_alpha=LORA_ALPHA,
    lora_dropout=0,
    bias="none",
    use_gradient_checkpointing="unsloth",
    random_state=3407,
    target_modules=[
        "q_proj", "k_proj", "v_proj", "o_proj",
        "gate_proj", "up_proj", "down_proj",
    ],
)

model.print_trainable_parameters()

gc.collect()
torch.cuda.empty_cache()
print(f"VRAM after LoRA: {torch.cuda.memory_allocated() / 1024**3:.1f} GB")
print("\n-> Run Cell 6 to set training arguments.")


# ============================================================
# CELL 6: Training Arguments
# ============================================================
from transformers import TrainingArguments

total_steps = len(dataset) // (BATCH_SIZE * GRAD_ACCUM) * EPOCHS
print(f"Estimated steps: {total_steps:,}")
print(f"Estimated time : ~{total_steps / 8 / 3600:.1f} hours\n")

args = TrainingArguments(
    output_dir=OUTPUT_DIR,
    overwrite_output_dir=True,
    num_train_epochs=EPOCHS,
    per_device_train_batch_size=BATCH_SIZE,
    gradient_accumulation_steps=GRAD_ACCUM,
    learning_rate=LR,
    lr_scheduler_type="cosine",
    warmup_steps=500,
    optim="adamw_8bit",
    weight_decay=0.01,
    max_grad_norm=0.3,
    bf16=torch.cuda.is_bf16_supported(),
    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",
    seed=3407,
)

print("Training arguments ready.")
print(f"  Effective batch: {BATCH_SIZE * GRAD_ACCUM}")
print(f"  LR: {LR}")
print(f"  Context: {MAX_SEQ_LENGTH}")
print(f"  Epochs: {EPOCHS}")
print("\n-> Run Cell 7 to initialize the trainer.")


# ============================================================
# CELL 7: Initialize SFTTrainer
# ============================================================
from trl import SFTTrainer

trainer = SFTTrainer(
    model=model,
    tokenizer=tokenizer,
    train_dataset=dataset,
    dataset_text_field="text",
    max_seq_length=MAX_SEQ_LENGTH,
    args=args,
    packing=False,
)

print("Trainer initialized.")
print(f"VRAM before training: {torch.cuda.memory_allocated() / 1024**3:.1f} GB")
print("\n-> Run Cell 8 to START TRAINING (takes hours).")


# ============================================================
# CELL 8: START TRAINING
# ============================================================
print(f"{'='*60}")
print("STARTING TRAINING")
print(f"{'='*60}")
print(f"  Conversations : {len(dataset):,}")
print(f"  Epochs        : {EPOCHS}")
print(f"  Batch         : {BATCH_SIZE} (eff: {BATCH_SIZE * GRAD_ACCUM})")
print(f"  Context       : {MAX_SEQ_LENGTH}")
print(f"  LoRA          : r={LORA_R}")
print(f"  LR            : {LR}")
print(f"  Expected time : ~{total_steps / 8 / 3600:.1f} hours")
print(f"{'='*60}")

trainer.train()

print("\nTraining complete!")
final_loss = trainer.state.log_history[-1].get('loss', 'N/A')
print(f"Final loss: {final_loss}")
print("\n-> Run Cell 9 to save the adapter.")


# ============================================================
# CELL 9: Save Adapter
# ============================================================
import os

os.makedirs(OUTPUT_DIR, exist_ok=True)

adapter_path = os.path.join(OUTPUT_DIR, "adapter")
model.save_pretrained(adapter_path)
tokenizer.save_pretrained(adapter_path)

size_mb = sum(
    os.path.getsize(os.path.join(adapter_path, f))
    for f in os.listdir(adapter_path)
) / (1024 * 1024)

print(f"Adapter saved: {adapter_path}")
print(f"Size: {size_mb:.1f} MB")

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: {merged_path}")
except RuntimeError as e:
    print(f"Merge skipped: {e}")

print("\n-> Run Cell 10 to test inference.")


# ============================================================
# CELL 10: Quick Test Inference
# ============================================================
messages = [
    {"role": "user", "content": "What does Charaka say about the qualities of a good physician in Sutrasthana?"}
]

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

print("Generating...")
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(f"\n{'='*60}")
print("RESPONSE")
print(f"{'='*60}")
print(response)
print(f"{'='*60}")