# ============================================================
# AYURVEDA GEMMA-4 TRAINING FOR A100 — FINAL CORRECTED VERSION
# All fixes applied: native messages, correct model, SFTConfig, MIG-safe LoRA
# ============================================================

# ============================================================
# 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

# --- USE HF MODEL (NOT GGUF) ---
MODEL_NAME = "google/gemma-4-12B-it"

# --- DATASET ---
DATASET_PATH = r"E:\dataset\gemma4_ayurveda_unsloth_clean_WITH_REFS.jsonl"
OUTPUT_DIR   = r"E:\dataset\gemma4_ayurveda_a100_output"

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:
    # MIG 20GB or similar conservative setup
    MAX_SEQ_LENGTH = 2048
    BATCH_SIZE = 1
    GRAD_ACCUM = 16
    LORA_R = 32
    LORA_ALPHA = 64
    LR = 2e-4
    print("\nDetected: ~20GB GPU -> MIG-safe 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
import gc

# Removed: from unsloth import standardize_sharegpt
# Reason: Dataset already uses native messages format (system/user/assistant)

dataset = load_dataset(
    "json",
    data_files=DATASET_PATH,
    split="train"
)

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

sample = dataset[0]
msgs = sample.get("messages", [])
roles = [m.get("role") for m in msgs]
print(f"Sample roles: {roles}")
print(f"Sample user : {msgs[1]['content'][:80]}...")

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


# ============================================================
# CELL 4: Load Gemma-4 12B Model (HF, not GGUF)
# ============================================================
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")
print(f"Chat template available: {tokenizer.chat_template is not None}")

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


# ============================================================
# CELL 5: Attach LoRA Adapter (MIG-safe r=32)
# ============================================================
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 format dataset with chat template.")


# ============================================================
# CELL 6: Format Dataset with Chat Template
# ============================================================
# Our dataset already has "messages" = [{role, content}, ...]
# We apply the tokenizer's chat template to create a "text" column for SFTTrainer
# Training data: add_generation_prompt=False (don't add model start token at end)

def format_chat_template(examples):
    texts = []
    for messages in examples["messages"]:
        text = tokenizer.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=False
        )
        texts.append(text)
    return {"text": texts}

dataset = dataset.map(format_chat_template, batched=True)

print(f"Formatted dataset. Sample text length: {len(dataset[0]['text'])} chars")
print(f"\n{'='*60}")
print("VERIFICATION — Chat template output (first 1000 chars):")
print(f"{'='*60}")
print(dataset[0]["text"][:1000])
print(f"\n{'='*60}")

# You should see:
# <start_of_turn>system
# You are an expert Ayurvedic scholar...
# <end_of_turn>
# <start_of_turn>user
# ...
# <start_of_turn>model
# ...
# <end_of_turn>

# Remove original messages column to save memory
dataset = dataset.remove_columns(["messages"])

gc.collect()
print("\n-> Run Cell 7 to set training arguments.")


# ============================================================
# CELL 7: Training Arguments (SFTConfig for new TRL)
# ============================================================
from trl import SFTConfig

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")

sft_args = SFTConfig(
    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=True,                 # A100 is Ampere, bf16 is native
    logging_steps=50,
    save_steps=500,
    save_total_limit=2,
    report_to="none",
    dataset_text_field="text",
    max_seq_length=MAX_SEQ_LENGTH,
    seed=3407,
)

print("SFTConfig 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 8 to initialize the trainer.")


# ============================================================
# CELL 8: Initialize SFTTrainer (new TRL syntax)
# ============================================================
from trl import SFTTrainer

trainer = SFTTrainer(
    model=model,
    train_dataset=dataset,
    processing_class=tokenizer,   # New TRL syntax (was tokenizer=)
    args=sft_args,
)

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


# ============================================================
# CELL 9: 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 10 to save the adapter.")


# ============================================================
# CELL 10: Save Adapter ONLY (skip 16-bit merge on MIG 20GB)
# ============================================================
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")
print("\nNOTE: 16-bit merge skipped for MIG 20GB safety.")
print("The LoRA adapter is your final output. Load it with:")
print(f"  model.load_adapter('{adapter_path}')")
print("\n-> Run Cell 11 to test inference.")


# ============================================================
# CELL 11: Quick Test Inference
# ============================================================
# For inference, add_generation_prompt=True so model knows to answer

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  # TRUE for inference
).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}")