{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {"id": "heading"},
   "source": [
    "# Brahmi Embedding Warmup — Aggressive Round 2\n",
    "\n",
    "**Goal**: Train the 300 new Devanagari embeddings using ALL 435K training rows.\n",
    "Round 1 (150 steps, 2000 chunks) produced garbled Hindi. This round uses:\n",
    "- Full corpus (435K rows → ~200K+ chunks)\n",
    "- Higher LR (1e-3)\n",
    "- Cosine schedule\n",
    "- Only chunks containing new tokens (efficient training)\n",
    "\n",
    "**Hardware**: T4 GPU (Colab free tier). ~2-3 hours.\n",
    "\n",
    "**Before starting**: Paste your HF_TOKEN below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {"id": "install"},
   "outputs": [],
   "source": [
    "!pip install -q transformers torch datasets accelerate huggingface_hub\n",
    "\n",
    "import json, os, math, time, random\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.utils.data import DataLoader\n",
    "from transformers import AutoTokenizer, AutoModelForCausalLM\n",
    "\n",
    "HF_TOKEN = \"hf_YOUR_TOKEN_HERE\"\n",
    "\n",
    "# Mount Google Drive for persistence across sessions\n",
    "from google.colab import drive\n",
    "drive.mount('/content/drive')\n",
    "DRIVE_DIR = '/content/drive/MyDrive/brahmi_training'\n",
    "os.makedirs(DRIVE_DIR, exist_ok=True)\n",
    "print(f'Saving to {DRIVE_DIR}')\n",
    "\n",
    "from huggingface_hub import login\n",
    "login(token=HF_TOKEN)\n",
    "\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "print(f'Device: {device}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {"id": "load"},
   "outputs": [],
   "source": [
    "MODEL_ID = 'eulogik/Bharat-Tiny-LLM-v2'\n",
    "\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",
    "\n",
    "model = AutoModelForCausalLM.from_pretrained(\n",
    "    MODEL_ID,\n",
    "    torch_dtype=torch.bfloat16,\n",
    "    device_map='auto',\n",
    "    token=HF_TOKEN,\n",
    ")\n",
    "\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",
    "print(f'Found {len(new_ids)} new tokens (IDs {new_ids[0]}–{new_ids[-1]})')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {"id": "freeze"},
   "outputs": [],
   "source": [
    "for p in model.parameters():\n",
    "    p.requires_grad = False\n",
    "model.model.embed_tokens.weight.requires_grad_(True)\n",
    "model.lm_head.weight.requires_grad_(True)\n",
    "\n",
    "class LearnableRows(nn.Module):\n",
    "    def __init__(self, embed, lmhead):\n",
    "        super().__init__()\n",
    "        self.embed = nn.Parameter(embed)\n",
    "        self.lmhead = nn.Parameter(lmhead)\n",
    "\n",
    "hidden = model.config.hidden_size\n",
    "embed_w = model.model.embed_tokens.weight\n",
    "lmhead_w = model.lm_head.weight\n",
    "\n",
    "learner = LearnableRows(\n",
    "    embed_w.data[new_ids].clone(),\n",
    "    lmhead_w.data[new_ids].clone()\n",
    ").to(device)\n",
    "\n",
    "n_params = len(new_ids) * hidden * 2\n",
    "print(f'Training {n_params:,} params (0.04% of {sum(p.numel() for p in model.parameters()):,})')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {"id": "data_load"},
   "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",
    "if not os.path.exists(DATA_FILE):\n",
    "    print('Downloading...')\n",
    "    r = requests.get(DATA_URL, stream=True,\n",
    "                     headers={'Authorization': f'Bearer {HF_TOKEN}'})\n",
    "    r.raise_for_status()\n",
    "    with open('/content/data.gz', 'wb') as f:\n",
    "        shutil.copyfileobj(r.raw, f)\n",
    "    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')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {"id": "tokenize"},
   "outputs": [],
   "source": [
    "CHUNKS_FILE = os.path.join(DRIVE_DIR, 'all_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",
    "else:\n",
    "    new_id_set = set(new_ids)\n",
    "    all_chunks = []\n",
    "    skipped = 0\n",
    "\n",
    "    for i, example in enumerate(dataset):\n",
    "        text = example.get('text', '')\n",
    "        if isinstance(text, str) and text.startswith('{'):\n",
    "            try:\n",
    "                data = json.loads(text)\n",
    "                text = ' '.join(m['content'] for m in data.get('messages', []))\n",
    "            except:\n",
    "                pass\n",
    "        ids = tok.encode(text)\n",
    "        for j in range(0, len(ids), 256):\n",
    "            chunk = ids[j:j+256]\n",
    "            if len(chunk) >= 10:\n",
    "                if any(tid in new_id_set for tid in chunk):\n",
    "                    all_chunks.append(chunk)\n",
    "                else:\n",
    "                    skipped += 1\n",
    "        if i % 50000 == 0 and i > 0:\n",
    "            print(f'  Processed {i}/{len(dataset)} rows...')\n",
    "\n",
    "    random.shuffle(all_chunks)\n",
    "    with open(CHUNKS_FILE, 'w') as f:\n",
    "        json.dump(all_chunks, f)\n",
    "    print(f'\\nKept {len(all_chunks)} chunks (skipped {skipped})')\n",
    "    print(f'Saved to Drive')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {"id": "dataloader"},
   "outputs": [],
   "source": [
    "def collate(batch):\n",
    "    mx = max(len(x) for x in batch)\n",
    "    pad = torch.zeros(len(batch), mx, dtype=torch.long)\n",
    "    for i, x in enumerate(batch):\n",
    "        pad[i, :len(x)] = torch.tensor(x, dtype=torch.long)\n",
    "    return pad.to(device)\n",
    "\n",
    "split = int(len(all_chunks) * 0.95)\n",
    "train_chunks = all_chunks[:split]\n",
    "val_chunks = all_chunks[split:]\n",
    "\n",
    "BATCH_SIZE = 4\n",
    "train_loader = DataLoader(train_chunks, batch_size=BATCH_SIZE, shuffle=True, collate_fn=collate)\n",
    "val_loader = DataLoader(val_chunks, batch_size=BATCH_SIZE, collate_fn=collate)\n",
    "\n",
    "print(f'Train: {len(train_chunks)} chunks ({len(train_loader)} batches)')\n",
    "print(f'Val: {len(val_chunks)} chunks ({len(val_loader)} batches)')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {"id": "init_val"},
   "outputs": [],
   "source": [
    "NUM_STEPS = 3000\n",
    "LR = 1e-3\n",
    "CKPT_DIR = os.path.join(DRIVE_DIR, 'checkpoints')\n",
    "os.makedirs(CKPT_DIR, exist_ok=True)\n",
    "\n",
    "optim = torch.optim.AdamW(learner.parameters(), lr=LR, weight_decay=0.01)\n",
    "scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(\n",
    "    optim, T_max=NUM_STEPS, eta_min=1e-5\n",
    ")\n",
    "\n",
    "# Resume from checkpoint if exists\n",
    "step = 0\n",
    "best_loss = float('inf')\n",
    "init_loss = None\n",
    "resume_path = os.path.join(CKPT_DIR, 'state.json')\n",
    "learner_path = os.path.join(CKPT_DIR, 'learner.pt')\n",
    "\n",
    "if os.path.exists(resume_path) and os.path.exists(learner_path):\n",
    "    with open(resume_path) as f:\n",
    "        state = json.load(f)\n",
    "    step = state['step']\n",
    "    best_loss = state['best_loss']\n",
    "    init_loss = state['init_loss']\n",
    "    learner.load_state_dict(torch.load(learner_path, map_location=device))\n",
    "    # Fast-forward scheduler\n",
    "    for _ in range(step):\n",
    "        scheduler.step()\n",
    "    print(f'Resumed from step {step}, best loss {best_loss:.4f}')\n",
    "else:\n",
    "    # Initial validation loss\n",
    "    model.eval()\n",
    "    init_losses = []\n",
    "    with torch.no_grad():\n",
    "        for i, batch in enumerate(val_loader):\n",
    "            init_losses.append(model(batch, labels=batch).loss.item())\n",
    "            if i >= 20: break\n",
    "    init_loss = np.mean(init_losses)\n",
    "    best_loss = init_loss\n",
    "    print(f'Initial val loss: {init_loss:.4f} (PPL: {math.exp(init_loss):.1f})')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {"id": "train_loop"},
   "outputs": [],
   "source": [
    "model.train()\n",
    "start_t = time.time()\n",
    "log_every = 100\n",
    "\n",
    "for epoch in range(10):\n",
    "    for batch in train_loader:\n",
    "        if step >= NUM_STEPS:\n",
    "            break\n",
    "\n",
    "        # Copy learner → model\n",
    "        with torch.no_grad():\n",
    "            embed_w.data[new_ids] = learner.embed.data.to(embed_w.device)\n",
    "            lmhead_w.data[new_ids] = learner.lmhead.data.to(lmhead_w.device)\n",
    "\n",
    "        out = model(batch, labels=batch)\n",
    "        loss = out.loss\n",
    "        loss.backward()\n",
    "\n",
    "        # Copy grads → learner\n",
    "        with torch.no_grad():\n",
    "            learner.embed.grad = embed_w.grad[new_ids].clone().to(learner.embed.device)\n",
    "            learner.lmhead.grad = lmhead_w.grad[new_ids].clone().to(learner.lmhead.device)\n",
    "\n",
    "        torch.nn.utils.clip_grad_norm_(learner.parameters(), 1.0)\n",
    "        optim.step()\n",
    "        scheduler.step()\n",
    "        optim.zero_grad()\n",
    "        model.zero_grad()\n",
    "        step += 1\n",
    "\n",
    "        if step % log_every == 0:\n",
    "            model.eval()\n",
    "            vlosses = []\n",
    "            with torch.no_grad():\n",
    "                for vb in val_loader:\n",
    "                    vlosses.append(model(vb, labels=vb).loss.item())\n",
    "                    if len(vlosses) >= 20: break\n",
    "            avg_vloss = np.mean(vlosses)\n",
    "            impr = (init_loss - avg_vloss) / init_loss * 100\n",
    "            lr_now = scheduler.get_last_lr()[0]\n",
    "            elapsed = time.time() - start_t\n",
    "            print(f'Step {step}/{NUM_STEPS} | '\n",
    "                  f'Tr:{loss.item():.4f} | '\n",
    "                  f'Val:{avg_vloss:.4f} (PPL:{math.exp(avg_vloss):.1f}) | '\n",
    "                  f'{impr:+.1f}% | '\n",
    "                  f'LR:{lr_now:.2e} | '\n",
    "                  f'{step/elapsed:.2f} it/s')\n",
    "            if avg_vloss < best_loss:\n",
    "                best_loss = avg_vloss\n",
    "                with torch.no_grad():\n",
    "                    embed_w.data[new_ids] = learner.embed.data.to(embed_w.device)\n",
    "                    lmhead_w.data[new_ids] = learner.lmhead.data.to(lmhead_w.device)\n",
    "                model.save_pretrained('brahmi-best')\n",
    "                tok.save_pretrained('brahmi-best')\n",
    "            # Save checkpoint for resume\n",
    "            torch.save(learner.state_dict(), learner_path)\n",
    "            with open(resume_path, 'w') as f:\n",
    "                json.dump({'step': step, 'best_loss': best_loss, 'init_loss': init_loss}, f)\n",
    "            model.train()\n",
    "\n",
    "    if step >= NUM_STEPS:\n",
    "        break\n",
    "\n",
    "elapsed = time.time() - start_t\n",
    "final_impr = (init_loss - best_loss) / init_loss * 100\n",
    "print(f'\\nDone! {step} steps in {elapsed:.0f}s ({step/elapsed:.2f} it/s)')\n",
    "print(f'Init loss: {init_loss:.4f} -> Best: {best_loss:.4f} ({final_impr:.1f}% improvement)')\n",
    "print(f'Init PPL: {math.exp(init_loss):.1f} -> Best PPL: {math.exp(best_loss):.1f}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {"id": "generate"},
   "outputs": [],
   "source": [
    "prompts = [\n",
    "    'मैं आपको बताना चाहता हूँ कि',\n",
    "    'भारत की राजधानी',\n",
    "    'नमस्ते, आप कैसे हैं? मैं',\n",
    "]\n",
    "\n",
    "for p in prompts:\n",
    "    inputs = tok(p, return_tensors='pt').to(model.device)\n",
    "    out = model.generate(\n",
    "        **inputs,\n",
    "        max_new_tokens=40,\n",
    "        temperature=0.3,\n",
    "        top_p=0.85,\n",
    "        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",
    "    print(f'Prompt: {p}')\n",
    "    print(f'  -> {gen}')\n",
    "    print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {"id": "upload"},
   "outputs": [],
   "source": [
    "from huggingface_hub import HfApi\n",
    "\n",
    "with torch.no_grad():\n",
    "    embed_w.data[new_ids] = learner.embed.data.to(embed_w.device)\n",
    "    lmhead_w.data[new_ids] = learner.lmhead.data.to(lmhead_w.device)\n",
    "\n",
    "model.save_pretrained('brahmi-trained')\n",
    "tok.save_pretrained('brahmi-trained')\n",
    "print('Saved locally')\n",
    "\n",
    "api = HfApi(token=HF_TOKEN)\n",
    "api.upload_folder(\n",
    "    folder_path='brahmi-trained',\n",
    "    repo_id='eulogik/Bharat-Tiny-LLM-v2',\n",
    "    repo_type='model',\n",
    "    commit_message=f'continued_pretrain_round2: {final_impr:.1f}% val loss improvement',\n",
    ")\n",
    "print('Uploaded to HF!')\n",
    "print('https://huggingface.co/eulogik/Bharat-Tiny-LLM-v2')"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {"provenance": []},
  "kernelspec": {"display_name": "Python 3", "name": "python3"},
  "language_info": {"name": "python"}
 },
 "nbformat": 4,
 "nbformat_minor": 0
}