{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Ayurveda Gemma-4 Fine-Tuning Script\n",
    "Simple Unsloth training for Charaka Samhita dataset.\n",
    "\n",
    "**Hardware**: RTX A4000 16GB (also works on bigger cards if you change model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ==========================================\n",
    "# 1. CONFIGURATION - Change these as needed\n",
    "# ==========================================\n",
    "\n",
    "# Pick your model. For A4000 16GB, use the 5B. For 24GB+ cards, try 12B.\n",
    "MODEL_NAME = \"unsloth/gemma-4-E2B-it\"          # ~5B params, fits 16GB\n",
    "# MODEL_NAME = \"unsloth/gemma-4-12b-it\"       # ~12B params, needs 24GB+\n",
    "\n",
    "# Dataset path (use the 1.69GB clean file for full training, or 404MB for quick tests)\n",
    "DATASET_PATH = r\"E:\\dataset\\gemma4_ayurveda_unsloth_clean.jsonl\"\n",
    "\n",
    "# Where to save the trained LoRA adapter\n",
    "OUTPUT_DIR = r\"E:\\dataset\\gemma4_ayurveda_lora_output\"\n",
    "\n",
    "# Training settings\n",
    "MAX_SEQ_LENGTH = 2048       # 2048 for 16GB safe. Use 3072/4096 if you have more VRAM.\n",
    "BATCH_SIZE = 1              # Keep at 1 for 16GB. Increase to 2 if you have 24GB+.\n",
    "GRAD_ACCUM = 8              # Effective batch = BATCH_SIZE * GRAD_ACCUM\n",
    "EPOCHS = 1                  # 1 epoch is enough for 971k conversations\n",
    "LORA_R = 64                 # Rank: 64 for deep domain (Sanskrit/medical). 32 for quick tests.\n",
    "LORA_ALPHA = 128            # Always 2x rank\n",
    "LR = 3e-4                   # Learning rate: 3e-4 for 5B, 2e-4 for 12B\n",
    "\n",
    "import os\n",
    "print(\"Dataset exists:\", os.path.exists(DATASET_PATH))\n",
    "print(\"Output dir:\", OUTPUT_DIR)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ==========================================\n",
    "# 2. INSTALL (if not already installed)\n",
    "# ==========================================\n",
    "\n",
    "# Uncomment and run this cell ONCE if you need to install:\n",
    "# !pip install unsloth transformers datasets trl accelerate bitsandbytes\n",
    "\n",
    "import torch\n",
    "from unsloth import FastLanguageModel, standardize_sharegpt\n",
    "from datasets import load_dataset\n",
    "from transformers import TrainingArguments\n",
    "from trl import SFTTrainer\n",
    "import gc\n",
    "\n",
    "print(\"CUDA available:\", torch.cuda.is_available())\n",
    "if torch.cuda.is_available():\n",
    "    print(\"GPU:\", torch.cuda.get_device_name(0))\n",
    "    print(\"VRAM total:\", f\"{torch.cuda.get_device_properties(0).total_memory / 1024**3:.1f} GB\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ==========================================\n",
    "# 3. LOAD DATASET\n",
    "# ==========================================\n",
    "\n",
    "dataset = load_dataset(\"json\", data_files=DATASET_PATH, split=\"train\")\n",
    "print(f\"Loaded {len(dataset):,} conversations\")\n",
    "\n",
    "# Convert ShareGPT format to what Unsloth expects\n",
    "dataset = standardize_sharegpt(dataset)\n",
    "print(\"Standardized. Sample text length:\", len(dataset[0][\"text\"]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ==========================================\n",
    "# 4. LOAD MODEL\n",
    "# ==========================================\n",
    "\n",
    "model, tokenizer = FastLanguageModel.from_pretrained(\n",
    "    model_name=MODEL_NAME,\n",
    "    max_seq_length=MAX_SEQ_LENGTH,\n",
    "    dtype=None,               # Auto-detect bf16 on Ampere, fp16 on older\n",
    "    load_in_4bit=True,        # QLoRA: 4-bit weights, huge VRAM saver\n",
    ")\n",
    "\n",
    "print(f\"Model loaded. VRAM used: {torch.cuda.memory_allocated() / 1024**3:.1f} GB\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ==========================================\n",
    "# 5. ATTACH LORA ADAPTER\n",
    "# ==========================================\n",
    "\n",
    "model = FastLanguageModel.get_peft_model(\n",
    "    model,\n",
    "    r=LORA_R,\n",
    "    lora_alpha=LORA_ALPHA,\n",
    "    lora_dropout=0,\n",
    "    bias=\"none\",\n",
    "    use_gradient_checkpointing=\"unsloth\",  # Saves ~30% VRAM\n",
    "    random_state=3407,\n",
    "    target_modules=[\n",
    "        \"q_proj\", \"k_proj\", \"v_proj\", \"o_proj\",\n",
    "        \"gate_proj\", \"up_proj\", \"down_proj\",\n",
    "    ],\n",
    ")\n",
    "\n",
    "# Print how many parameters are trainable\n",
    "model.print_trainable_parameters()\n",
    "\n",
    "gc.collect()\n",
    "torch.cuda.empty_cache()\n",
    "print(f\"VRAM after LoRA: {torch.cuda.memory_allocated() / 1024**3:.1f} GB\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ==========================================\n",
    "# 6. TRAINING ARGUMENTS\n",
    "# ==========================================\n",
    "\n",
    "total_steps = len(dataset) // (BATCH_SIZE * GRAD_ACCUM) * EPOCHS\n",
    "print(f\"Estimated training steps: {total_steps:,}\")\n",
    "print(f\"Estimated time on A4000: ~{total_steps / 2.5 / 3600:.1f} hours\")\n",
    "\n",
    "args = TrainingArguments(\n",
    "    output_dir=OUTPUT_DIR,\n",
    "    overwrite_output_dir=True,\n",
    "\n",
    "    num_train_epochs=EPOCHS,\n",
    "    per_device_train_batch_size=BATCH_SIZE,\n",
    "    gradient_accumulation_steps=GRAD_ACCUM,\n",
    "\n",
    "    learning_rate=LR,\n",
    "    lr_scheduler_type=\"cosine\",\n",
    "    warmup_steps=500,\n",
    "\n",
    "    optim=\"adamw_8bit\",\n",
    "    weight_decay=0.01,\n",
    "    max_grad_norm=0.3,\n",
    "\n",
    "    bf16=torch.cuda.is_bf16_supported(),\n",
    "    fp16=not torch.cuda.is_bf16_supported(),\n",
    "\n",
    "    logging_steps=50,\n",
    "    save_strategy=\"steps\",\n",
    "    save_steps=2000,\n",
    "    save_total_limit=2,\n",
    "\n",
    "    group_by_length=True,\n",
    "    report_to=\"none\",  # Change to \"wandb\" if you want logging\n",
    "    seed=3407,\n",
    ")\n",
    "\n",
    "print(\"Training args ready.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ==========================================\n",
    "# 7. TRAINER\n",
    "# ==========================================\n",
    "\n",
    "trainer = SFTTrainer(\n",
    "    model=model,\n",
    "    tokenizer=tokenizer,\n",
    "    train_dataset=dataset,\n",
    "    dataset_text_field=\"text\",\n",
    "    max_seq_length=MAX_SEQ_LENGTH,\n",
    "    args=args,\n",
    "    packing=False,  # NEVER pack for instruction/chat data\n",
    ")\n",
    "\n",
    "print(\"Trainer ready. Run next cell to start training!\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ==========================================\n",
    "# 8. TRAIN!\n",
    "# ==========================================\n",
    "\n",
    "trainer.train()\n",
    "print(\"\\nTraining complete!\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ==========================================\n",
    "# 9. SAVE ADAPTER\n",
    "# ==========================================\n",
    "\n",
    "import os\n",
    "os.makedirs(OUTPUT_DIR, exist_ok=True)\n",
    "\n",
    "adapter_path = os.path.join(OUTPUT_DIR, \"adapter\")\n",
    "model.save_pretrained(adapter_path)\n",
    "tokenizer.save_pretrained(adapter_path)\n",
    "\n",
    "print(f\"Adapter saved to: {adapter_path}\")\n",
    "print(\"Size:\", f\"{sum(os.path.getsize(os.path.join(adapter_path, f)) for f in os.listdir(adapter_path)) / 1024**2:.1f} MB\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ==========================================\n",
    "# 10. OPTIONAL: Merge LoRA into Base Model\n",
    "# ==========================================\n",
    "# Only run this if you have enough VRAM (needs ~12GB for 5B, ~24GB for 12B)\n",
    "# Creates a standalone model that doesn't need the adapter.\n",
    "\n",
    "try:\n",
    "    merged_path = os.path.join(OUTPUT_DIR, \"merged_16bit\")\n",
    "    model.save_pretrained_merged(merged_path, tokenizer, save_method=\"merged_16bit\")\n",
    "    print(f\"Merged model saved to: {merged_path}\")\n",
    "except RuntimeError as e:\n",
    "    print(f\"Merge skipped (OOM or error): {e}\")\n",
    "    print(\"You can still use the adapter + base model for inference.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ==========================================\n",
    "# 11. QUICK TEST INFERENCE\n",
    "# ==========================================\n",
    "\n",
    "messages = [\n",
    "    {\"role\": \"user\", \"content\": \"What does Charaka say about the qualities of a good physician in Sutrasthana?\"}\n",
    "]\n",
    "\n",
    "inputs = tokenizer.apply_chat_template(\n",
    "    messages, \n",
    "    tokenize=True, \n",
    "    return_tensors=\"pt\", \n",
    "    add_generation_prompt=True\n",
    ").to(\"cuda\")\n",
    "\n",
    "outputs = model.generate(\n",
    "    inputs, \n",
    "    max_new_tokens=512, \n",
    "    temperature=0.7, \n",
    "    top_p=0.9, \n",
    "    do_sample=True\n",
    ")\n",
    "\n",
    "print(tokenizer.decode(outputs[0], skip_special_tokens=True))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.11.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}