{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "f4d27c45-ea3c-4833-b83c-65766d112aea",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "GPUs visible : 1\n",
      "GPU 0        : NVIDIA A100-SXM4-40GB MIG 3g.20gb\n",
      "VRAM         : 19.5 GB\n",
      "\n",
      "Detected MIG 20GB -> Conservative settings\n",
      "\n",
      "==================================================\n",
      "CONFIG\n",
      "==================================================\n",
      "Model      : unsloth/gemma-3-4b-it\n",
      "Dataset    : /nlsasfs/home/aikosh/prod-aikosh35/gemma4_ayurveda_unsloth_clean_WITH_REFS.jsonl\n",
      "Output     : /nlsasfs/home/aikosh/prod-aikosh35/gemma3_ayurveda_a100_output\n",
      "Seq Length : 1024\n",
      "Batch      : 1 (eff: 32)\n",
      "LoRA       : r=16, alpha=16\n",
      "LR         : 0.0002\n",
      "Epochs     : 1\n",
      "==================================================\n",
      "\n",
      "Dataset found: 1722.2 MB\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import os\n",
    "\n",
    "# FIX UNSLOTH CACHE PERMISSIONS\n",
    "os.environ[\"UNSLOTH_COMPILE_CACHE_DIR\"] = (\n",
    "    \"/nlsasfs/home/aikosh/prod-aikosh35/unsloth_cache\"\n",
    ")\n",
    "os.makedirs(os.environ[\"UNSLOTH_COMPILE_CACHE_DIR\"], exist_ok=True)\n",
    "\n",
    "# --- GEMMA 3 MODEL ---\n",
    "MODEL_NAME = \"unsloth/gemma-3-4b-it\"\n",
    "\n",
    "DATASET_PATH = \"/nlsasfs/home/aikosh/prod-aikosh35/gemma4_ayurveda_unsloth_clean_WITH_REFS.jsonl\"\n",
    "OUTPUT_DIR   = \"/nlsasfs/home/aikosh/prod-aikosh35/gemma3_ayurveda_a100_output\"\n",
    "\n",
    "\n",
    "if not torch.cuda.is_available():\n",
    "    raise RuntimeError(\"No GPU found\")\n",
    "\n",
    "gpu_count = torch.cuda.device_count()\n",
    "vram_gb = torch.cuda.get_device_properties(0).total_memory / 1024**3\n",
    "\n",
    "print(f\"GPUs visible : {gpu_count}\")\n",
    "print(f\"GPU 0        : {torch.cuda.get_device_name(0)}\")\n",
    "print(f\"VRAM         : {vram_gb:.1f} GB\")\n",
    "\n",
    "\n",
    "MAX_SEQ_LENGTH = 1024\n",
    "BATCH_SIZE = 1\n",
    "GRAD_ACCUM = 32\n",
    "LORA_R = 16\n",
    "LORA_ALPHA = 16\n",
    "LR = 2e-4\n",
    "EPOCHS = 1\n",
    "\n",
    "\n",
    "print(\"\\nDetected MIG 20GB -> Conservative settings\")\n",
    "\n",
    "print(f\"\"\"\n",
    "==================================================\n",
    "CONFIG\n",
    "==================================================\n",
    "Model      : {MODEL_NAME}\n",
    "Dataset    : {DATASET_PATH}\n",
    "Output     : {OUTPUT_DIR}\n",
    "Seq Length : {MAX_SEQ_LENGTH}\n",
    "Batch      : {BATCH_SIZE} (eff: {BATCH_SIZE*GRAD_ACCUM})\n",
    "LoRA       : r={LORA_R}, alpha={LORA_ALPHA}\n",
    "LR         : {LR}\n",
    "Epochs     : {EPOCHS}\n",
    "==================================================\n",
    "\"\"\")\n",
    "\n",
    "\n",
    "if os.path.exists(DATASET_PATH):\n",
    "    print(f\"Dataset found: {os.path.getsize(DATASET_PATH)/1024**2:.1f} MB\")\n",
    "else:\n",
    "    print(\"Dataset missing\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "a16b7df2-92f2-4bae-8844-26d2f030f03b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loaded 971,590 conversations\n",
      "['system', 'user', 'assistant']\n",
      "What significance did the Ayurvedic sage attach to longevity when establishing the framework for Sut\n",
      "Run Cell 4\n"
     ]
    }
   ],
   "source": [
    "from datasets import load_dataset\n",
    "import gc\n",
    "\n",
    "dataset = load_dataset(\n",
    "    \"json\",\n",
    "    data_files=DATASET_PATH,\n",
    "    split=\"train\"\n",
    ")\n",
    "\n",
    "print(f\"Loaded {len(dataset):,} conversations\")\n",
    "\n",
    "sample = dataset[0]\n",
    "msgs = sample[\"messages\"]\n",
    "\n",
    "print([x[\"role\"] for x in msgs])\n",
    "print(msgs[1][\"content\"][:100])\n",
    "\n",
    "gc.collect()\n",
    "\n",
    "print(\"Run Cell 4\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "d9fe66e3-9f4b-4aa0-ac9f-6b6e58a96f3f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Cache: /nlsasfs/home/aikosh/prod-aikosh35/.cache/unsloth\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "\n",
    "CACHE = \"/nlsasfs/home/aikosh/prod-aikosh35/.cache/unsloth\"\n",
    "\n",
    "os.makedirs(CACHE, exist_ok=True)\n",
    "os.makedirs(f\"{CACHE}/torch\", exist_ok=True)\n",
    "os.makedirs(f\"{CACHE}/triton\", exist_ok=True)\n",
    "\n",
    "os.environ[\"UNSLOTH_COMPILED_CACHE\"] = CACHE\n",
    "os.environ[\"UNSLOTH_CACHE_DIR\"] = CACHE\n",
    "os.environ[\"TORCHINDUCTOR_CACHE_DIR\"] = f\"{CACHE}/torch\"\n",
    "os.environ[\"TRITON_CACHE_DIR\"] = f\"{CACHE}/triton\"\n",
    "os.environ[\"XDG_CACHE_HOME\"] = CACHE\n",
    "\n",
    "print(\"Cache:\", CACHE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "7bf796c7-bbb0-48f1-b9d1-bd86c7f4b102",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "🦥 Unsloth: Will patch your computer to enable 2x faster free finetuning.\n",
      "🦥 Unsloth Zoo will now patch everything to make training faster!\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[unsloth_zoo.log|WARNING]Loading unsloth/gemma-3-4b-it as text-only; vision/audio towers skipped. Use FastVisionModel for multimodal inputs.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Unsloth: WARNING `trust_remote_code` is True.\n",
      "Are you certain you want to do remote code execution?\n",
      "==((====))==  Unsloth 2026.6.8: Fast Gemma3 patching. Transformers: 5.5.0.\n",
      "   \\\\   /|    NVIDIA A100-SXM4-40GB MIG 3g.20gb. Num GPUs = 1. Max memory: 19.5 GB. Platform: Linux.\n",
      "O^O/ \\_/ \\    Torch: 2.10.0+cu128. CUDA: 8.0. CUDA Toolkit: 12.8. Triton: 3.6.0\n",
      "\\        /    Bfloat16 = TRUE. FA [Xformers = 0.0.35. FA2 = False]\n",
      " \"-____-\"     Free license: http://github.com/unslothai/unsloth\n",
      "Unsloth: Fast downloading is enabled - ignore downloading bars which are red colored!\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "691a8f21117646dea5d5a217a3b8562a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading weights:   0%|          | 0/444 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[1mGemma3ForCausalLM LOAD REPORT\u001b[0m from: unsloth/gemma-3-4b-it-unsloth-bnb-4bit\n",
      "Key                                                                         | Status     |  | \n",
      "----------------------------------------------------------------------------+------------+--+-\n",
      "vision_tower.vision_model.encoder.layers.{0...26}.mlp.fc2.weight            | UNEXPECTED |  | \n",
      "vision_tower.vision_model.encoder.layers.{0...26}.self_attn.k_proj.bias     | UNEXPECTED |  | \n",
      "vision_tower.vision_model.encoder.layers.{0...26}.self_attn.q_proj.bias     | UNEXPECTED |  | \n",
      "vision_tower.vision_model.encoder.layers.{0...26}.layer_norm2.bias          | UNEXPECTED |  | \n",
      "vision_tower.vision_model.encoder.layers.{0...26}.self_attn.k_proj.weight   | UNEXPECTED |  | \n",
      "vision_tower.vision_model.encoder.layers.{0...26}.mlp.fc1.bias              | UNEXPECTED |  | \n",
      "vision_tower.vision_model.post_layernorm.weight                             | UNEXPECTED |  | \n",
      "vision_tower.vision_model.encoder.layers.{0...26}.self_attn.q_proj.weight   | UNEXPECTED |  | \n",
      "vision_tower.vision_model.encoder.layers.{0...26}.layer_norm1.weight        | UNEXPECTED |  | \n",
      "vision_tower.vision_model.encoder.layers.{0...26}.self_attn.v_proj.weight   | UNEXPECTED |  | \n",
      "vision_tower.vision_model.encoder.layers.{0...26}.self_attn.out_proj.bias   | UNEXPECTED |  | \n",
      "vision_tower.vision_model.encoder.layers.{0...26}.layer_norm1.bias          | UNEXPECTED |  | \n",
      "vision_tower.vision_model.encoder.layers.{0...26}.mlp.fc1.weight            | UNEXPECTED |  | \n",
      "vision_tower.vision_model.encoder.layers.{0...26}.self_attn.out_proj.weight | UNEXPECTED |  | \n",
      "vision_tower.vision_model.encoder.layers.{0...26}.mlp.fc2.bias              | UNEXPECTED |  | \n",
      "vision_tower.vision_model.encoder.layers.{0...26}.self_attn.v_proj.bias     | UNEXPECTED |  | \n",
      "vision_tower.vision_model.encoder.layers.{0...26}.layer_norm2.weight        | UNEXPECTED |  | \n",
      "vision_tower.vision_model.embeddings.patch_embedding.bias                   | UNEXPECTED |  | \n",
      "vision_tower.vision_model.post_layernorm.bias                               | UNEXPECTED |  | \n",
      "multi_modal_projector.mm_input_projection_weight                            | UNEXPECTED |  | \n",
      "vision_tower.vision_model.embeddings.patch_embedding.weight                 | UNEXPECTED |  | \n",
      "multi_modal_projector.mm_soft_emb_norm.weight                               | UNEXPECTED |  | \n",
      "vision_tower.vision_model.embeddings.position_embedding.weight              | UNEXPECTED |  | \n",
      "\n",
      "Notes:\n",
      "- UNEXPECTED:\tcan be ignored when loading from different task/architecture; not ok if you expect identical arch.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model loaded\n",
      "VRAM: 3.5 GB\n"
     ]
    }
   ],
   "source": [
    "from unsloth import FastLanguageModel\n",
    "import torch\n",
    "import gc\n",
    "\n",
    "model, tokenizer = FastLanguageModel.from_pretrained(\n",
    "    model_name=\"unsloth/gemma-3-4b-it\",\n",
    "    max_seq_length=MAX_SEQ_LENGTH,\n",
    "    dtype=None,\n",
    "    load_in_4bit=True,\n",
    "    trust_remote_code=True,\n",
    "    text_only=True,\n",
    ")\n",
    "\n",
    "print(\"Model loaded\")\n",
    "print(f\"VRAM: {torch.cuda.memory_allocated()/1024**3:.1f} GB\")\n",
    "\n",
    "gc.collect()\n",
    "torch.cuda.empty_cache()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "5319c9f1-5d38-494a-934a-929c7e15bf66",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Unsloth: Explicit target_modules are constrained by the finetune_(vision|language|attention|mlp) filters; adapters attach only where both select.\n",
      "trainable params: 29,802,496 || all params: 3,910,065,664 || trainable%: 0.7622\n",
      "VRAM after LoRA: 3.6 GB\n"
     ]
    }
   ],
   "source": [
    "from unsloth import FastLanguageModel\n",
    "import torch, gc\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\",\n",
    "    random_state=3407,\n",
    "\n",
    "    target_modules=[\n",
    "        \"q_proj\",\n",
    "        \"k_proj\",\n",
    "        \"v_proj\",\n",
    "        \"o_proj\",\n",
    "        \"gate_proj\",\n",
    "        \"up_proj\",\n",
    "        \"down_proj\",\n",
    "    ],\n",
    ")\n",
    "\n",
    "model.print_trainable_parameters()\n",
    "\n",
    "model.config.use_cache = False\n",
    "\n",
    "gc.collect()\n",
    "torch.cuda.empty_cache()\n",
    "\n",
    "print(f\"VRAM after LoRA: {torch.cuda.memory_allocated()/1024**3:.1f} GB\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "35561388-43ca-4c8c-b8d8-2dca66631713",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Formatted dataset: 971,590\n",
      "<bos><start_of_turn>user\n",
      "You are an expert Ayurvedic scholar trained in Charaka Samhita. Provide accurate responses with precise citations including full Sanskrit section names from Sutrasthana (सूत्रस्थान).\n",
      "\n",
      "What significance did the Ayurvedic sage attach to longevity when establishing the framework for Sutrasthana (सूत्रस्थान)?<end_of_turn>\n",
      "<start_of_turn>model\n",
      "Now (I) shall expound the chapter on longevity. [1]\n",
      "\n",
      "This teaching from **Sutrasthana (सूत्रस्थान)** establishes foundational principles for General Principles and Fundamentals. Located in Chapter 1, Verse 1, this verse contributes to the systematic understanding of Ayurvedic medicine.\n",
      "\n",
      "**Sanskrit:** अथातो दीर्घञ्जीवितीयमध्यायं व्याख्यास्यामः ||१||\n",
      "\n",
      "**Transliteration:** athāto dīrghañjīvitīyamadhyāyaṃ vyākhyāsyāmaḥ ||1||\n",
      "\n",
      "**Full Reference:** Sutrasthana (सूत्रस्थान) Chapter 1, Verse 1\n",
      "\n",
      "**Reference ID:** CH_SU_1/1<end_of_turn>\n",
      "\n",
      "\n",
      "-> Run Cell 7\n"
     ]
    }
   ],
   "source": [
    "# CELL 6: Format Dataset with Chat Template\n",
    "\n",
    "def format_chat_template(examples):\n",
    "    texts = []\n",
    "    for messages in examples[\"messages\"]:\n",
    "        text = tokenizer.apply_chat_template(\n",
    "            messages,\n",
    "            tokenize=False,\n",
    "            add_generation_prompt=False\n",
    "        )\n",
    "        texts.append(text)\n",
    "    return {\"text\": texts}\n",
    "\n",
    "dataset = dataset.map(\n",
    "    format_chat_template,\n",
    "    batched=True,\n",
    "    remove_columns=[\"messages\"]\n",
    ")\n",
    "\n",
    "print(f\"Formatted dataset: {len(dataset):,}\")\n",
    "print(dataset[0][\"text\"][:1000])\n",
    "\n",
    "gc.collect()\n",
    "torch.cuda.empty_cache()\n",
    "\n",
    "print(\"\\n-> Run Cell 7\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "b8a02632-ec5e-4463-8981-20fc91ed1331",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SFTConfig ready\n"
     ]
    }
   ],
   "source": [
    "# CELL 7: Training Args (Gemma3 / TRL 2026)\n",
    "\n",
    "from trl import SFTConfig\n",
    "\n",
    "total_steps = len(dataset) // (BATCH_SIZE * GRAD_ACCUM) * EPOCHS\n",
    "\n",
    "sft_args = SFTConfig(\n",
    "    output_dir=OUTPUT_DIR,\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=100,\n",
    "\n",
    "    optim=\"adamw_8bit\",\n",
    "    weight_decay=0.01,\n",
    "    max_grad_norm=0.3,\n",
    "\n",
    "    bf16=True,\n",
    "\n",
    "    logging_steps=25,\n",
    "    save_steps=250,\n",
    "    save_total_limit=2,\n",
    "\n",
    "    report_to=\"none\",\n",
    "\n",
    "    dataset_text_field=\"text\",\n",
    "    max_length=MAX_SEQ_LENGTH,\n",
    "\n",
    "    gradient_checkpointing=True,\n",
    "    seed=3407,\n",
    ")\n",
    "\n",
    "print(\"SFTConfig ready\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "8bf71e9e-d848-42c9-a211-667ba764250b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b6740b6a713445f6b90cfc72f71ae2f4",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Adding EOS to train dataset:   0%|          | 0/971590 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "183daed3521543fdb0899238109327f5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Tokenizing train dataset:   0%|          | 0/971590 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "441fb52147e847ceaea616dfb1a86b74",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Truncating train dataset:   0%|          | 0/971590 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Trainer ready\n",
      "VRAM: 3.6 GB\n",
      "Run Cell 9\n"
     ]
    }
   ],
   "source": [
    "# CELL 8: Trainer\n",
    "\n",
    "from trl import SFTTrainer\n",
    "\n",
    "trainer = SFTTrainer(\n",
    "    model=model,\n",
    "    train_dataset=dataset,\n",
    "    processing_class=tokenizer,\n",
    "    args=sft_args,\n",
    ")\n",
    "\n",
    "print(\"Trainer ready\")\n",
    "print(f\"VRAM: {torch.cuda.memory_allocated()/1024**3:.1f} GB\")\n",
    "\n",
    "print(\"Run Cell 9\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "2f3dcd3c-0f4b-49cb-beb7-d5dea33a16ca",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ[\"PYTORCH_CUDA_ALLOC_CONF\"] = \"expandable_segments:True\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "3cd28459-84a7-4ff3-9e72-31f6eb657d77",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "============================================================\n",
      "STARTING GEMMA-3 4B QLoRA TRAINING\n",
      "============================================================\n"
     ]
    },
    {
     "ename": "NameError",
     "evalue": "name 'train_dataset' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[9], line 10\u001b[0m\n\u001b[1;32m      8\u001b[0m \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mSTARTING GEMMA-3 4B QLoRA TRAINING\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m      9\u001b[0m \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m=\u001b[39m\u001b[38;5;124m\"\u001b[39m \u001b[38;5;241m*\u001b[39m \u001b[38;5;241m60\u001b[39m)\n\u001b[0;32m---> 10\u001b[0m \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mTrain set : \u001b[39m\u001b[38;5;132;01m{\u001b[39;00m\u001b[38;5;28mlen\u001b[39m(\u001b[43mtrain_dataset\u001b[49m)\u001b[38;5;132;01m:\u001b[39;00m\u001b[38;5;124m,\u001b[39m\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m     11\u001b[0m \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mEval set  : \u001b[39m\u001b[38;5;132;01m{\u001b[39;00m\u001b[38;5;28mlen\u001b[39m(eval_dataset)\u001b[38;5;132;01m:\u001b[39;00m\u001b[38;5;124m,\u001b[39m\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m     12\u001b[0m \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mEpochs    : \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mEPOCHS\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m)\n",
      "\u001b[0;31mNameError\u001b[0m: name 'train_dataset' is not defined"
     ]
    }
   ],
   "source": [
    "# ============================================================\n",
    "# CELL 9 — FULL TRAINING\n",
    "# ============================================================\n",
    "gc.collect()\n",
    "torch.cuda.empty_cache()\n",
    "\n",
    "print(\"=\" * 60)\n",
    "print(\"STARTING GEMMA-3 4B QLoRA TRAINING\")\n",
    "print(\"=\" * 60)\n",
    "print(f\"Train set : {len(train_dataset):,}\")\n",
    "print(f\"Eval set  : {len(eval_dataset):,}\")\n",
    "print(f\"Epochs    : {EPOCHS}\")\n",
    "print(f\"Batch     : {BATCH_SIZE}  (accum: {GRAD_ACCUM}, effective: {BATCH_SIZE * GRAD_ACCUM})\")\n",
    "print(f\"Seq Len   : {MAX_SEQ_LENGTH}\")\n",
    "print(f\"LoRA      : r={LORA_R}, alpha={LORA_ALPHA}\")\n",
    "print(f\"LR        : {LR}\")\n",
    "print(f\"VRAM      : {torch.cuda.memory_allocated() / 1024**3:.1f} GB\")\n",
    "print(\"=\" * 60)\n",
    "\n",
    "# resume_from_checkpoint=True  →  starts from latest checkpoint if one exists in OUTPUT_DIR\n",
    "#                                 safe for fresh runs (starts from scratch if none found)\n",
    "trainer.train(resume_from_checkpoint=True)\n",
    "\n",
    "print(\"=\" * 60)\n",
    "print(\"TRAINING COMPLETE\")\n",
    "print(\"=\" * 60)\n",
    "print(\"Cell 9 complete → Run Cell 10\")\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "1032bcba-9b6a-44ce-a509-8c207f28717f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2.10.0+cu128\n",
      "/nlsasfs/home/aikosh/prod-aikosh35/.venv/lib/python3.10/site-packages/torch/__init__.py\n",
      "🦥 Unsloth: Will patch your computer to enable 2x faster free finetuning.\n",
      "🦥 Unsloth Zoo will now patch everything to make training faster!\n",
      "SUCCESS\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "print(torch.__version__)\n",
    "print(torch.__file__)\n",
    "\n",
    "from unsloth import FastLanguageModel\n",
    "print(\"SUCCESS\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e346ab6f-0a24-420e-b017-b32cc41d884e",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
