Complete rewrite: single setup cell, all imports in cell 1

brahmi_lora_qlora.ipynb

@@ -2,79 +2,60 @@
  "cells": [
   {
    "cell_type": "markdown",
-   "metadata": {
-    "id": "heading"
-   },
+   "metadata": {"id": "heading"},
    "source": [
-    "# Bharat-Tiny-LLM: LoRA/QLoRA Fine-tuning\n",
+    "# Bharat-Tiny-LLM v2: QLoRA Fine-tuning\n",
     "\n",
-    "**Goal**: Fine-tune the full Qwen2.5-1.5B model with LoRA adapters to improve\n",
-    "Hindi generation quality with the 300 new Devanagari tokens.\n",
+    "**Goal**: Fine-tune Qwen2.5-1.5B (with 300 new Devanagari tokens) using QLoRA.\n",
+    "This teaches the model to USE the new tokens for coherent Hindi generation.\n",
     "\n",
-    "**Approach**: QLoRA (4-bit quantized LoRA) for memory efficiency on T4 GPU.\n",
-    "Trains attention layers while keeping base weights frozen.\n",
+    "**Hardware**: T4 GPU (Colab free). ~1 hour for 20K chunks, 2 epochs.\n",
     "\n",
-    "**Hardware**: T4 GPU (Colab free tier). ~2-3 hours.\n",
-    "\n",
-    "**Before starting**: Paste your HF_TOKEN below."
+    "**DO NOT run cells manually** — use Runtime → Run all."
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "id": "install"
-   },
+   "metadata": {"id": "setup"},
    "outputs": [],
    "source": [
-    "import os\n",
-    "from transformers import TrainingArguments\n",
+    "!pip install -q transformers torch datasets accelerate huggingface_hub peft bitsandbytes\n",
+    "\n",
+    "import json, os, math, time, random, requests, gzip, shutil\n",
+    "import numpy as np\n",
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "from torch.utils.data import DataLoader, Dataset\n",
+    "from transformers import (\n",
+    "    AutoTokenizer, AutoModelForCausalLM,\n",
+    "    BitsAndBytesConfig, TrainingArguments, Trainer\n",
+    ")\n",
+    "from peft import LoraConfig, get_peft_model, TaskType, prepare_model_for_kbit_training\n",
+    "from huggingface_hub import login, HfApi\n",
+    "from google.colab import drive\n",
+    "from datasets import load_dataset\n",
     "\n",
-    "DRIVE_DIR = '/content/drive/MyDrive/brahmi_lora'\n",
-    "CKPT_DIR = os.path.join(DRIVE_DIR, 'checkpoints')\n",
-    "os.makedirs(CKPT_DIR, exist_ok=True)\n",
+    "# === CONFIG ===\n",
+    "HF_TOKEN = \"hf_YOUR_TOKEN_HERE\"\n",
+    "MODEL_ID = 'eulogik/Bharat-Tiny-LLM-v2'\n",
     "\n",
-    "training_args = TrainingArguments(\n",
-    "    output_dir=CKPT_DIR,\n",
-    "    num_train_epochs=2,\n",
-    "    per_device_train_batch_size=4,\n",
-    "    per_device_eval_batch_size=4,\n",
-    "    gradient_accumulation_steps=4,\n",
-    "    learning_rate=2e-4,\n",
-    "    weight_decay=0.01,\n",
-    "    warmup_ratio=0.1,\n",
-    "    lr_scheduler_type='cosine',\n",
-    "    logging_steps=25,\n",
-    "    eval_strategy='steps',\n",
-    "    eval_steps=100,\n",
-    "    save_strategy='steps',\n",
-    "    save_steps=100,\n",
-    "    save_total_limit=3,\n",
-    "    load_best_model_at_end=True,\n",
-    "    metric_for_best_model='eval_loss',\n",
-    "    greater_is_better=False,\n",
-    "    bf16=True,\n",
-    "    gradient_checkpointing=True,\n",
-    "    gradient_checkpointing_kwargs={'use_reentrant': False},\n",
-    "    optim='paged_adamw_8bit',\n",
-    "    max_grad_norm=1.0,\n",
-    "    report_to='none',\n",
-    ")\n",
+    "drive.mount('/content/drive')\n",
+    "DRIVE_DIR = '/content/drive/MyDrive/brahmi_lora'\n",
+    "os.makedirs(DRIVE_DIR, exist_ok=True)\n",
     "\n",
-    "print(f'Effective batch size: {training_args.per_device_train_batch_size * training_args.gradient_accumulation_steps}')"
+    "login(token=HF_TOKEN)\n",
+    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
+    "print(f'Device: {device}')\n",
+    "print(f'Drive: {DRIVE_DIR}')"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "id": "load_model"
-   },
+   "metadata": {"id": "load_model"},
    "outputs": [],
    "source": [
-    "MODEL_ID = 'eulogik/Bharat-Tiny-LLM-v2'\n",
-    "\n",
-    "# 4-bit quantization config (QLoRA)\n",
     "bnb_config = BitsAndBytesConfig(\n",
     "    load_in_4bit=True,\n",
     "    bnb_4bit_quant_type='nf4',\n",
@@ -82,7 +63,6 @@
     "    bnb_4bit_use_double_quant=True,\n",
     ")\n",
     "\n",
-    "print(f'Loading {MODEL_ID} in 4-bit...')\n",
     "tok = AutoTokenizer.from_pretrained(MODEL_ID, token=HF_TOKEN)\n",
     "if tok.pad_token is None:\n",
     "    tok.pad_token = tok.eos_token\n",
@@ -96,25 +76,20 @@
     "\n",
     "model = prepare_model_for_kbit_training(model)\n",
     "\n",
-    "# Identify new tokens\n",
     "new_ids = sorted(\n",
     "    tid for tid, t in tok.added_tokens_decoder.items()\n",
     "    if not str(t).startswith('<') and not getattr(t, 'special', False)\n",
     ")\n",
-    "new_id_set = set(new_ids)\n",
     "print(f'Found {len(new_ids)} new tokens')\n",
-    "print(f'Model loaded: {sum(p.numel() for p in model.parameters()) / 1e6:.1f}M params')"
+    "print(f'Params: {sum(p.numel() for p in model.parameters()) / 1e6:.1f}M')"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "id": "lora_config"
-   },
+   "metadata": {"id": "lora"},
    "outputs": [],
    "source": [
-    "# LoRA configuration\n",
     "lora_config = LoraConfig(\n",
     "    task_type=TaskType.CAUSAL_LM,\n",
     "    r=16,\n",
@@ -125,20 +100,15 @@
     ")\n",
     "\n",
     "model = get_peft_model(model, lora_config)\n",
-    "model.print_trainable_parameters()\n",
-    "\n"
+    "model.print_trainable_parameters()"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "id": "data_load"
-   },
+   "metadata": {"id": "download_data"},
    "outputs": [],
    "source": [
-    "import requests, gzip, shutil\n",
-    "\n",
     "DATA_URL = 'https://huggingface.co/eulogik/Bharat-Tiny-LLM-v2/resolve/main/train_gold_v3.jsonl.gz'\n",
     "DATA_FILE = '/content/train_gold_v3.jsonl'\n",
     "\n",
@@ -152,9 +122,7 @@
     "    with gzip.open('/content/data.gz', 'rb') as gz, open(DATA_FILE, 'wb') as f:\n",
     "        shutil.copyfileobj(gz, f)\n",
     "    os.remove('/content/data.gz')\n",
-    "    print(f'Ready! {os.path.getsize(DATA_FILE) / 1e6:.0f} MB')\n",
     "\n",
-    "from datasets import load_dataset\n",
     "dataset = load_dataset('text', data_files=DATA_FILE, split='train')\n",
     "print(f'Loaded {len(dataset)} rows')"
    ]
@@ -162,19 +130,15 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "id": "tokenize"
-   },
+   "metadata": {"id": "tokenize"},
    "outputs": [],
    "source": [
     "CHUNKS_FILE = os.path.join(DRIVE_DIR, 'lora_chunks.json')\n",
     "\n",
     "if os.path.exists(CHUNKS_FILE):\n",
-    "    print('Loading cached chunks from Drive...')\n",
     "    with open(CHUNKS_FILE) as f:\n",
     "        all_chunks = json.load(f)\n",
-    "    random.shuffle(all_chunks)\n",
-    "    print(f'Loaded {len(all_chunks)} chunks')\n",
+    "    print(f'Loaded {len(all_chunks)} chunks from Drive')\n",
     "else:\n",
     "    all_chunks = []\n",
     "    for i, example in enumerate(dataset):\n",
@@ -194,12 +158,11 @@
     "            print(f'  Processed {i}/{len(dataset)} rows...')\n",
     "\n",
     "    random.shuffle(all_chunks)\n",
-    "    # Cap at 20K chunks for feasible T4 training\n",
     "    if len(all_chunks) > 20000:\n",
     "        all_chunks = all_chunks[:20000]\n",
     "    with open(CHUNKS_FILE, 'w') as f:\n",
     "        json.dump(all_chunks, f)\n",
-    "    print(f'\\nTotal chunks: {len(all_chunks)}')\n",
+    "    print(f'Total chunks: {len(all_chunks)} (saved to Drive)')\n",
     "\n",
     "split = int(len(all_chunks) * 0.95)\n",
     "train_chunks = all_chunks[:split]\n",
@@ -210,12 +173,10 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "id": "dataset_class"
-   },
+   "metadata": {"id": "dataset"},
    "outputs": [],
    "source": [
-    "class ChunkDataset(torch.utils.data.Dataset):\n",
+    "class ChunkDataset(Dataset):\n",
     "    def __init__(self, chunks, max_len=512):\n",
     "        self.chunks = chunks\n",
     "        self.max_len = max_len\n",
@@ -225,26 +186,23 @@
     "\n",
     "    def __getitem__(self, idx):\n",
     "        ids = self.chunks[idx][:self.max_len]\n",
-    "        # Pad to max_len\n",
     "        padded = ids + [tok.pad_token_id] * (self.max_len - len(ids))\n",
-    "        attention_mask = [1] * len(ids) + [0] * (self.max_len - len(ids))\n",
+    "        mask = [1] * len(ids) + [0] * (self.max_len - len(ids))\n",
     "        return {\n",
     "            'input_ids': torch.tensor(padded, dtype=torch.long),\n",
-    "            'attention_mask': torch.tensor(attention_mask, dtype=torch.long),\n",
+    "            'attention_mask': torch.tensor(mask, dtype=torch.long),\n",
     "            'labels': torch.tensor(padded, dtype=torch.long),\n",
     "        }\n",
     "\n",
     "train_dataset = ChunkDataset(train_chunks)\n",
     "val_dataset = ChunkDataset(val_chunks)\n",
-    "print(f'Train dataset: {len(train_dataset)} samples')"
+    "print(f'Train: {len(train_dataset)}, Val: {len(val_dataset)}')"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "id": "training_args"
-   },
+   "metadata": {"id": "train"},
    "outputs": [],
    "source": [
     "CKPT_DIR = os.path.join(DRIVE_DIR, 'checkpoints')\n",
@@ -252,19 +210,19 @@
     "\n",
     "training_args = TrainingArguments(\n",
     "    output_dir=CKPT_DIR,\n",
-    "    num_train_epochs=3,\n",
-    "    per_device_train_batch_size=2,\n",
-    "    per_device_eval_batch_size=2,\n",
-    "    gradient_accumulation_steps=8,\n",
+    "    num_train_epochs=2,\n",
+    "    per_device_train_batch_size=4,\n",
+    "    per_device_eval_batch_size=4,\n",
+    "    gradient_accumulation_steps=4,\n",
     "    learning_rate=2e-4,\n",
     "    weight_decay=0.01,\n",
     "    warmup_ratio=0.1,\n",
     "    lr_scheduler_type='cosine',\n",
-    "    logging_steps=50,\n",
+    "    logging_steps=25,\n",
     "    eval_strategy='steps',\n",
-    "    eval_steps=200,\n",
+    "    eval_steps=100,\n",
     "    save_strategy='steps',\n",
-    "    save_steps=200,\n",
+    "    save_steps=100,\n",
     "    save_total_limit=3,\n",
     "    load_best_model_at_end=True,\n",
     "    metric_for_best_model='eval_loss',\n",
@@ -277,19 +235,6 @@
     "    report_to='none',\n",
     ")\n",
     "\n",
-    "print(f'Effective batch size: {training_args.per_device_train_batch_size * training_args.gradient_accumulation_steps}')"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "id": "train"
-   },
-   "outputs": [],
-   "source": [
-    "from transformers import Trainer\n",
-    "\n",
     "trainer = Trainer(\n",
     "    model=model,\n",
     "    args=training_args,\n",
@@ -297,86 +242,71 @@
     "    eval_dataset=val_dataset,\n",
     ")\n",
     "\n",
-    "# Resume from checkpoint if exists\n",
-    "resume_ckpt = None\n",
-    "CKPT_DIR = '/content/drive/MyDrive/brahmi_lora/checkpoints'\n",
-    "if os.path.exists(CKPT_DIR):\n",
-    "    checkpoints = [d for d in os.listdir(CKPT_DIR) if d.startswith('checkpoint-')]\n",
-    "    if checkpoints:\n",
-    "        latest = max(checkpoints, key=lambda x: int(x.split('-')[1]))\n",
-    "        resume_ckpt = os.path.join(CKPT_DIR, latest)\n",
-    "        print(f'Resuming from {latest}')\n",
+    "# Resume from last checkpoint if exists\n",
+    "checkpoints = [d for d in os.listdir(CKPT_DIR) if d.startswith('checkpoint-')]\n",
+    "resume = None\n",
+    "if checkpoints:\n",
+    "    latest = max(checkpoints, key=lambda x: int(x.split('-')[1]))\n",
+    "    resume = os.path.join(CKPT_DIR, latest)\n",
+    "    print(f'Resuming from {latest}')\n",
     "\n",
+    "stats = training_args.per_device_train_batch_size * training_args.gradient_accumulation_steps\n",
+    "print(f'Effective batch size: {stats}')\n",
     "print('Starting training...')\n",
-    "trainer.train(resume_from_checkpoint=resume_ckpt)"
+    "trainer.train(resume_from_checkpoint=resume)"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "id": "save_upload"
-   },
+   "metadata": {"id": "save_upload"},
    "outputs": [],
    "source": [
-    "import os\n",
-    "\n",
-    "DRIVE_DIR = '/content/drive/MyDrive/brahmi_lora'\n",
+    "# Save LoRA adapter to Drive and HF\n",
     "ADAPTER_DIR = os.path.join(DRIVE_DIR, 'lora_adapter')\n",
     "model.save_pretrained(ADAPTER_DIR)\n",
     "tok.save_pretrained(ADAPTER_DIR)\n",
-    "print(f'LoRA adapter saved to {ADAPTER_DIR}')\n",
+    "print(f'LoRA saved to {ADAPTER_DIR}')\n",
     "\n",
-    "from huggingface_hub import HfApi\n",
     "api = HfApi(token=HF_TOKEN)\n",
-    "\n",
     "api.upload_folder(\n",
     "    folder_path=ADAPTER_DIR,\n",
     "    repo_id='eulogik/Bharat-Tiny-LLM-v2-LoRA',\n",
     "    repo_type='model',\n",
-    "    commit_message='LoRA adapter: 3 epochs, 2e-4 LR, rank=16',\n",
+    "    commit_message='LoRA: 2 epochs, rank=16, 2e-4',\n",
     ")\n",
-    "print('Uploaded to HF!')\n",
-    "print('https://huggingface.co/eulogik/Bharat-Tiny-LLM-v2-LoRA')"
+    "print('Uploaded LoRA to HF!')"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "id": "merge_upload"
-   },
+   "metadata": {"id": "merge"},
    "outputs": [],
    "source": [
-    "import os\n",
-    "\n",
-    "DRIVE_DIR = '/content/drive/MyDrive/brahmi_lora'\n",
+    "# Merge LoRA into base and upload\n",
     "print('Merging LoRA into base model...')\n",
-    "merged_model = model.merge_and_unload()\n",
+    "merged = model.merge_and_unload()\n",
     "\n",
-    "MERGED_DIR = os.path.join(DRIVE_DIR, 'merged_model')\n",
-    "merged_model.save_pretrained(MERGED_DIR)\n",
+    "MERGED_DIR = os.path.join(DRIVE_DIR, 'merged')\n",
+    "merged.save_pretrained(MERGED_DIR)\n",
     "tok.save_pretrained(MERGED_DIR)\n",
     "print(f'Merged model saved to {MERGED_DIR}')\n",
     "\n",
-    "from huggingface_hub import HfApi\n",
     "api = HfApi(token=HF_TOKEN)\n",
     "api.upload_folder(\n",
     "    folder_path=MERGED_DIR,\n",
     "    repo_id='eulogik/Bharat-Tiny-LLM-v2',\n",
     "    repo_type='model',\n",
-    "    commit_message='LoRA merged: improved Hindi generation',\n",
+    "    commit_message='LoRA merged: rank=16, 2 epochs',\n",
     ")\n",
-    "print('Merged model uploaded!')\n",
-    "print('https://huggingface.co/eulogik/Bharat-Tiny-LLM-v2')"
+    "print('Merged model uploaded to HF!')"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "id": "generate"
-   },
+   "metadata": {"id": "generate"},
    "outputs": [],
    "source": [
     "prompts = [\n",
@@ -395,27 +325,18 @@
     "        repetition_penalty=1.25,\n",
     "        do_sample=True,\n",
     "    )\n",
-    "    gen = tok.decode(out[0][inputs.input_ids.shape[-1]:],\n",
-    "                     skip_special_tokens=True)\n",
+    "    gen = tok.decode(out[0][inputs.input_ids.shape[-1]:], skip_special_tokens=True)\n",
     "    print(f'Prompt: {p}')\n",
-    "    print(f'  -> {gen}')\n",
-    "    print()"
+    "    print(f'  -> {gen}\\n')"
    ]
   }
  ],
  "metadata": {
   "accelerator": "GPU",
-  "colab": {
-   "provenance": []
-  },
-  "kernelspec": {
-   "display_name": "Python 3",
-   "name": "python3"
-  },
-  "language_info": {
-   "name": "python"
-  }
+  "colab": {"provenance": []},
+  "kernelspec": {"display_name": "Python 3", "name": "python3"},
+  "language_info": {"name": "python"}
  },
  "nbformat": 4,
  "nbformat_minor": 0
-}
+}