refactor: simplify notebook with single ROOT_DIR variable

- Uses ROOT_DIR as single source of truth for all paths
- All paths use os.path.join() with absolute paths
- Fresh clone every time to avoid cached issues
- Searches for data in multiple possible locations

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

colab_train_stack29.ipynb

@@ -8,22 +8,14 @@
     "\n",
     "**Zero-cost training on Google Colab free tier with T4 GPU**\n",
     "\n",
-    "This notebook trains a LoRA adapter for Stack 2.9 on **Qwen2.5-Coder-7B** using a free T4 GPU.\n",
-    "\n",
     "⏱️ **Expected runtime:** 3-5 hours\n",
     "💾 **VRAM needed:** ~12GB (fits in T4's 15GB)\n",
-    "📦 **Output:** `./training_output/` (on Google Drive)\n",
     "\n",
     "---\n",
     "\n",
-    "**CRITICAL:** All data saved to **Google Drive** to persist through disconnects.\n",
-    "\n",
-    "**Instructions:**\n",
-    "1. Runtime → Change runtime type → **GPU (T4)**\n",
-    "2. Run all cells in order\n",
-    "3. **Allow** Drive access when prompted\n",
+    "**CRITICAL:** Run cells in order from the top!\n",
     "\n",
-    "---\n"
+    "---"
    ]
   },
   {
@@ -32,19 +24,17 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Check GPU\n",
-    "!nvidia-smi"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## 1️⃣ Mount Google Drive (REQUIRED for persistence)\n",
+    "# STEP 1: Setup - Mount Drive and define root directory\n",
+    "from google.colab import drive\n",
+    "drive.mount('/content/drive')\n",
     "\n",
-    "Click the link, allow access, copy the auth code, paste it, and press Enter.\n",
+    "import os\n",
+    "ROOT_DIR = \"/content/drive/MyDrive/stack-2.9\"\n",
+    "os.makedirs(ROOT_DIR, exist_ok=True)\n",
+    "os.chdir(ROOT_DIR)\n",
     "\n",
-    "**Without Drive mounting, training will be lost if Colab disconnects!**"
+    "print(f\"✅ Working directory: {os.getcwd()}\")\n",
+    "!ls -la"
    ]
   },
   {
@@ -52,29 +42,19 @@
    "execution_count": null,
    "metadata": {},
    "outputs": [],
-   "source": "from google.colab import drive\ndrive.mount('/content/drive')\n\n# Set up paths on Drive - ALL OUTPUT GOES HERE\nimport os\nBASE_PATH = \"/content/drive/MyDrive/stack-2.9-colab\"\nos.makedirs(BASE_PATH, exist_ok=True)\nos.chdir(BASE_PATH)\nprint(f\"\\n✅ Working directory: {os.getcwd()}\")\nprint(f\"All outputs will be saved to: {BASE_PATH}\")\nprint(\"\\nCurrent folder contents:\")\n!ls -la"
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## 2️⃣ Clone Stack 2.9 Repository"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": "# Remove old clone if exists and re-clone fresh\nimport os\nimport shutil\n\nif os.path.exists('stack-2.9'):\n    print(\"⚠️ Removing old stack-2.9 directory...\")\n    shutil.rmtree('stack-2.9')\n\n!git clone https://github.com/my-ai-stack/stack-2.9.git\n\nos.chdir('stack-2.9')\nprint(f\"Now in: {os.getcwd()}\")\n!ls -la"
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
    "source": [
-    "## 3️⃣ Install Dependencies\n",
+    "# STEP 2: Clone repo (fresh every time)\n",
+    "import shutil\n",
     "\n",
-    "Takes 5-10 minutes."
+    "if os.path.exists('stack-2.9'):\n",
+    "    print(\"Removing old stack-2.9...\")\n",
+    "    shutil.rmtree('stack-2.9')\n",
+    "\n",
+    "!git clone https://github.com/my-ai-stack/stack-2.9.git\n",
+    "\n",
+    "os.chdir(os.path.join(ROOT_DIR, 'stack-2.9'))\n",
+    "print(f\"✅ In: {os.getcwd()}\")\n",
+    "!ls -la"
    ]
   },
   {
@@ -82,15 +62,11 @@
    "execution_count": null,
    "metadata": {},
    "outputs": [],
-   "source": "!pip install --upgrade pip\n!pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118\n!pip install transformers==4.40.0 peft==0.10.0 accelerate datasets pyyaml tqdm scipy\n!pip install bitsandbytes==0.43.3\nprint(\"\\n✅ Dependencies installed\")"
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
    "source": [
-    "## 4️⃣ Download Base Model (Qwen2.5-Coder-7B)\n",
-    "\n",
-    "This takes ~15-20 minutes. The model is ~14GB."
+    "# STEP 3: Install dependencies\n",
+    "!pip install -q torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118\n",
+    "!pip install -q transformers peft accelerate datasets pyyaml tqdm scipy bitsandbytes\n",
+    "print(\"✅ Dependencies installed\")"
    ]
   },
   {
@@ -98,15 +74,25 @@
    "execution_count": null,
    "metadata": {},
    "outputs": [],
-   "source": "MODEL_NAME = \"Qwen/Qwen2.5-Coder-7B\"\n\n# Use absolute path for model (relative to current working directory)\nimport os\nMODEL_DIR = os.path.abspath(\"./base_model_qwen7b\")\n\nfrom transformers import AutoModelForCausalLM, AutoTokenizer\n\nif not os.path.exists(MODEL_DIR):\n    print(f\"Downloading {MODEL_NAME} to {MODEL_DIR}...\")\n    print(\"This will take 15-20 minutes...\")\n    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)\n    tokenizer.save_pretrained(MODEL_DIR)\n    model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, trust_remote_code=True)\n    model.save_pretrained(MODEL_DIR)\n    print(f\"✅ Model saved to {MODEL_DIR}\")\nelse:\n    print(f\"✅ Model already exists at {MODEL_DIR}\")\n\nprint(f\"\\n📁 Model absolute path: {MODEL_DIR}\")\n!ls -lh {MODEL_DIR} | head -10"
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
    "source": [
-    "## 5️⃣ Prepare Training Data\n",
+    "# STEP 4: Download Base Model (Qwen2.5-Coder-7B)\n",
+    "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
+    "\n",
+    "MODEL_NAME = \"Qwen/Qwen2.5-Coder-7B\"\n",
+    "MODEL_DIR = os.path.join(ROOT_DIR, \"stack-2.9/base_model_qwen7b\")\n",
+    "\n",
+    "if not os.path.exists(os.path.join(MODEL_DIR, \"config.json\")):\n",
+    "    print(f\"Downloading {MODEL_NAME} to {MODEL_DIR}...\")\n",
+    "    print(\"This will take 15-20 minutes...\")\n",
+    "    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)\n",
+    "    tokenizer.save_pretrained(MODEL_DIR)\n",
+    "    model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, trust_remote_code=True)\n",
+    "    model.save_pretrained(MODEL_DIR)\n",
+    "    print(f\"✅ Model saved\")\n",
+    "else:\n",
+    "    print(f\"✅ Model already exists\")\n",
     "\n",
-    "Using the bundled mini dataset (5K examples) for quick prototyping."
+    "!ls -lh {MODEL_DIR} | head -5"
    ]
   },
   {
@@ -114,13 +100,28 @@
    "execution_count": null,
    "metadata": {},
    "outputs": [],
-   "source": "# Check if data exists in the repo\nimport os\n\n# First check if data directory exists in repo\nrepo_data_path = os.path.join(os.getcwd(), \"data/final/train.jsonl\")\ndata_alt_path = os.path.join(os.getcwd(), \"training-data/final/train.jsonl\")\n\nif os.path.exists(repo_data_path):\n    DATA_PATH = os.path.abspath(repo_data_path)\n    print(f\"✅ Training data found at {DATA_PATH}\")\n    !wc -l {DATA_PATH}\nelif os.path.exists(data_alt_path):\n    DATA_PATH = os.path.abspath(data_alt_path)\n    print(f\"✅ Training data found at {DATA_PATH}\")\n    !wc -l {DATA_PATH}\nelse:\n    print(\"⚠️ Data not found in repo. Checking what's available:\")\n    !find . -name \"*.jsonl\" 2>/dev/null | head -10\n    \n    # If still no data, use a fallback - create small test dataset\n    print(\"\\n⚠️ Creating small test dataset (500 examples) for testing...\")\n    !python scripts/create_mini_dataset.py --size 500 --output data_mini/train_mini.jsonl --source ./data/final/train.jsonl 2>/dev/null || echo \"Script failed\"\n    DATA_PATH = os.path.abspath(\"./data_mini/train_mini.jsonl\")\n    if os.path.exists(DATA_PATH):\n        !ls -lh {DATA_PATH}\n    else:\n        raise FileNotFoundError(\"Could not create or find training data\")\n\nprint(f\"\\n📁 Data absolute path: {DATA_PATH}\")"
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
    "source": [
-    "## 6️⃣ Prepare Training Configuration"
+    "# STEP 5: Find training data\n",
+    "REPO_DIR = os.path.join(ROOT_DIR, \"stack-2.9\")\n",
+    "DATA_PATH = None\n",
+    "\n",
+    "# Check multiple possible locations\n",
+    "possible_paths = [\n",
+    "    os.path.join(REPO_DIR, \"data/final/train.jsonl\"),\n",
+    "    os.path.join(REPO_DIR, \"training-data/final/train.jsonl\"),\n",
+    "    os.path.join(REPO_DIR, \"data_mini/train_mini.jsonl\"),\n",
+    "]\n",
+    "\n",
+    "for path in possible_paths:\n",
+    "    if os.path.exists(path):\n",
+    "        DATA_PATH = path\n",
+    "        print(f\"✅ Found data at: {path}\")\n",
+    "        break\n",
+    "\n",
+    "if DATA_PATH is None:\n",
+    "    print(\"❌ No training data found!\")\n",
+    "    print(\"\\nSearching for jsonl files:\")\n",
+    "    !find {REPO_DIR} -name \"*.jsonl\" | head -10"
    ]
   },
   {
@@ -128,35 +129,37 @@
    "execution_count": null,
    "metadata": {},
    "outputs": [],
-   "source": "# Use Colab config and update paths\nimport yaml\nimport os\n\n# Use absolute path for config\nREPO_DIR = os.getcwd()\nconfig_path = os.path.join(REPO_DIR, \"stack/training/train_config_local.yaml\")\n\n# Check if config exists\nif not os.path.exists(config_path):\n    print(f\"❌ Config not found at: {config_path}\")\n    print(\"📁 Checking repo structure:\")\n    !find . -name \"train_config*.yaml\" | head -10\n    raise FileNotFoundError(f\"Config file not found: {config_path}\")\n\nprint(f\"📄 Loading config from: {config_path}\")\nwith open(config_path, 'r') as f:\n    config = yaml.safe_load(f)\n\n# Update for Colab/T4 GPU - use absolute paths\nconfig['model']['name'] = MODEL_DIR\nconfig['data']['input_path'] = DATA_PATH\nconfig['output']['lora_dir'] = os.path.abspath(\"./training_output/lora\")\nconfig['output']['merged_dir'] = os.path.abspath(\"./training_output/merged\")\nconfig['hardware']['device'] = \"cuda\"  # Use T4 GPU\nconfig['hardware']['num_gpus'] = 1\n\n# Save updated config\nOUTPUT_DIR = os.path.abspath(\"./training_output\")\nos.makedirs(OUTPUT_DIR, exist_ok=True)\nupdated_config_path = f\"{OUTPUT_DIR}/train_config.yaml\"\n\nwith open(updated_config_path, 'w') as f:\n    yaml.dump(config, f)\n\nprint(f\"✅ Config saved to: {updated_config_path}\")\nprint(\"\\nConfig summary:\")\nprint(f\"  - Model: {config['model']['name']}\")\nprint(f\"  - Data: {config['data']['input_path']}\")\nprint(f\"  - Device: {config['hardware']['device']}\")\nprint(f\"  - Epochs: {config['training']['num_epochs']}\")"
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
    "source": [
-    "## 7️⃣ Train LoRA Adapter\n",
+    "# STEP 6: Prepare Training Configuration\n",
+    "import yaml\n",
     "\n",
-    "⚠️ **This takes 3-5 hours. DO NOT INTERRUPT.**\n",
+    "config_path = os.path.join(REPO_DIR, \"stack/training/train_config_local.yaml\")\n",
     "\n",
-    "If Colab disconnects, reconnect and training will resume from checkpoint automatically.\n",
+    "if not os.path.exists(config_path):\n",
+    "    raise FileNotFoundError(f\"Config not found at: {config_path}\")\n",
     "\n",
-    "Watch for `Train loss:` decreasing. It should start ~2.0-3.0 and trend downward.\n",
+    "with open(config_path, 'r') as f:\n",
+    "    config = yaml.safe_load(f)\n",
     "\n",
-    "Checkpoints saved every 200 steps to `./training_output/lora/` (on Drive)."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": "import os\nimport sys\n\n# Add training module to path\nsys.path.insert(0, './stack/training')\n\nprint(\"=\"*60)\nprint(\"STARTING TRAINING\")\nprint(\"=\"*60)\nprint(f\"Working directory: {os.getcwd()}\")\nprint(f\"Config: {updated_config_path}\")\nprint(f\"Output: {OUTPUT_DIR}/lora\")\nprint(\"=\"*60 + \"\\n\")\n\n# Run training\nfrom train_lora import train_lora\n\ntrainer = train_lora(updated_config_path)\n\nprint(\"\\n\" + \"=\"*60)\nprint(\"TRAINING FINISHED OR STOPPED\")\nprint(\"=\"*60)"
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## 8️⃣ Verify Training Output"
+    "# Update config with absolute paths\n",
+    "config['model']['name'] = MODEL_DIR\n",
+    "config['data']['input_path'] = DATA_PATH\n",
+    "OUTPUT_DIR = os.path.join(ROOT_DIR, \"training_output\")\n",
+    "config['output']['lora_dir'] = os.path.join(OUTPUT_DIR, \"lora\")\n",
+    "config['output']['merged_dir'] = os.path.join(OUTPUT_DIR, \"merged\")\n",
+    "config['hardware']['device'] = \"cuda\"\n",
+    "config['hardware']['num_gpus'] = 1\n",
+    "\n",
+    "os.makedirs(OUTPUT_DIR, exist_ok=True)\n",
+    "updated_config_path = os.path.join(OUTPUT_DIR, \"train_config.yaml\")\n",
+    "\n",
+    "with open(updated_config_path, 'w') as f:\n",
+    "    yaml.dump(config, f)\n",
+    "\n",
+    "print(f\"✅ Config saved to: {updated_config_path}\")\n",
+    "print(f\"   Model: {config['model']['name']}\")\n",
+    "print(f\"   Data: {config['data']['input_path']}\")\n",
+    "print(f\"   Device: {config['hardware']['device']}\")"
    ]
   },
   {
@@ -165,26 +168,20 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "lora_dir = f\"{OUTPUT_DIR}/lora\"\n",
-    "print(f\"Checking LoRA output: {lora_dir}\")\n",
+    "# STEP 7: Train LoRA Adapter\n",
+    "import sys\n",
+    "sys.path.insert(0, os.path.join(REPO_DIR, \"stack/training\"))\n",
     "\n",
-    "if os.path.exists(lora_dir):\n",
-    "    files = os.listdir(lora_dir)\n",
-    "    print(f\"✅ LoRA adapter found! {len(files)} files:\")\n",
-    "    for f in files:\n",
-    "        size = os.path.getsize(os.path.join(lora_dir, f)) / (1024*1024)\n",
-    "        print(f\"   - {f}: {size:.1f} MB\")\n",
-    "else:\n",
-    "    print(\"⚠️ LoRA directory not found - training may have failed\")"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## 9️⃣ Merge LoRA Adapter with Base Model\n",
+    "print(\"=\"*60)\n",
+    "print(\"STARTING TRAINING\")\n",
+    "print(\"=\"*60)\n",
     "\n",
-    "Combines the trained adapter with the base model to produce a standalone fine-tuned model."
+    "from train_lora import train_lora\n",
+    "trainer = train_lora(updated_config_path)\n",
+    "\n",
+    "print(\"=\"*60)\n",
+    "print(\"TRAINING COMPLETED\")\n",
+    "print(\"=\"*60)"
    ]
   },
   {
@@ -192,13 +189,14 @@
    "execution_count": null,
    "metadata": {},
    "outputs": [],
-   "source": "import yaml\nimport sys\nimport os\n\nsys.path.insert(0, './stack/training')\nfrom merge_adapter import merge_adapter\n\nmerged_dir = os.path.abspath(f\"{OUTPUT_DIR}/merged\")\n\nprint(\"=\"*60)\nprint(\"MERGING LORA ADAPTER\")\nprint(\"=\"*60)\nprint(f\"Base model: {MODEL_DIR}\")\nprint(f\"LoRA path: {OUTPUT_DIR}/lora\")\nprint(f\"Output path: {merged_dir}\")\nprint(\"=\"*60)\n\n# Create merge config\nmerge_config = {\n    'model': {'name': MODEL_DIR, 'trust_remote_code': True},\n    'output': {'lora_dir': f'{OUTPUT_DIR}/lora', 'merged_dir': merged_dir},\n    'quantization': {'enabled': False}\n}\n\nmerge_config_path = f\"{OUTPUT_DIR}/merge_config.yaml\"\nwith open(merge_config_path, 'w') as f:\n    yaml.dump(merge_config, f)\n\n# Run merge\nmerge_adapter(\n    config_path=merge_config_path,\n    lora_path=f\"{OUTPUT_DIR}/lora\",\n    output_path=merged_dir\n)\n\nprint(f\"\\n✅ Merged model saved to: {merged_dir}\")\n!ls -lh {merged_dir}/"
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
    "source": [
-    "## 🔟 Test Inference (Quick Check)"
+    "# STEP 8: Verify and Merge\n",
+    "lora_dir = os.path.join(OUTPUT_DIR, \"lora\")\n",
+    "print(f\"Checking LoRA: {lora_dir}\")\n",
+    "if os.path.exists(lora_dir):\n",
+    "    !ls -lh {lora_dir}\n",
+    "else:\n",
+    "    print(\"❌ No LoRA output found\")"
    ]
   },
   {
@@ -207,42 +205,27 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from transformers import AutoTokenizer, AutoModelForCausalLM\n",
-    "import torch\n",
-    "\n",
-    "model_path = f\"{OUTPUT_DIR}/merged\"\n",
-    "print(f\"Loading model from {model_path}...\")\n",
-    "\n",
-    "try:\n",
-    "    tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)\n",
-    "    model = AutoModelForCausalLM.from_pretrained(\n",
-    "        model_path,\n",
-    "        torch_dtype=torch.float16,\n",
-    "        device_map=\"cuda\",\n",
-    "        trust_remote_code=True\n",
-    "    )\n",
-    "    \n",
-    "    prompt = \"Write a Python function to reverse a string:\\n\\n```python\\n\"\n",
-    "    inputs = tokenizer(prompt, return_tensors=\"pt\").to(\"cuda\")\n",
-    "    \n",
-    "    print(\"Generating...\")\n",
-    "    with torch.no_grad():\n",
-    "        outputs = model.generate(\n",
-    "            **inputs,\n",
-    "            max_new_tokens=200,\n",
-    "            temperature=0.2,\n",
-    "            do_sample=True,\n",
-    "            pad_token_id=tokenizer.eos_token_id\n",
-    "        )\n",
-    "    \n",
-    "    response = tokenizer.decode(outputs[0], skip_special_tokens=True)\n",
-    "    print(\"=\"*40)\n",
-    "    print(\"RESPONSE:\")\n",
-    "    print(\"=\"*40)\n",
-    "    print(response[len(prompt):])\n",
-    "except Exception as e:\n",
-    "    print(f\"❌ Error: {e}\")\n",
-    "    print(\"\\nThis is expected if training hasn't completed yet.\")"
+    "# STEP 9: Merge LoRA\n",
+    "import sys\n",
+    "sys.path.insert(0, os.path.join(REPO_DIR, \"stack/training\"))\n",
+    "from merge_adapter import merge_adapter\n",
+    "\n",
+    "merged_dir = os.path.join(OUTPUT_DIR, \"merged\")\n",
+    "os.makedirs(merged_dir, exist_ok=True)\n",
+    "\n",
+    "merge_config = {\n",
+    "    'model': {'name': MODEL_DIR, 'trust_remote_code': True},\n",
+    "    'output': {'lora_dir': lora_dir, 'merged_dir': merged_dir},\n",
+    "    'quantization': {'enabled': False}\n",
+    "}\n",
+    "\n",
+    "merge_config_path = os.path.join(OUTPUT_DIR, \"merge_config.yaml\")\n",
+    "with open(merge_config_path, 'w') as f:\n",
+    "    yaml.dump(merge_config, f)\n",
+    "\n",
+    "merge_adapter(merge_config_path, lora_dir, merged_dir)\n",
+    "print(f\"✅ Merged to: {merged_dir}\")\n",
+    "!ls -lh {merged_dir}"
    ]
   },
   {
@@ -251,22 +234,17 @@
    "source": [
     "## 🔚 Training Complete!\n",
     "\n",
-    "Your model is ready in `./training_output/merged/` and saved to Google Drive.\n",
-    "\n",
-    "**Next steps:**\n",
-    "1. **Download** `training_output/merged/` from Drive to your local machine\n",
-    "2. **Run evaluation**: Evaluate on HumanEval/MBPP benchmarks\n",
-    "3. **Upload model** to Hugging Face Hub\n",
-    "4. **Apply to Together AI**\n",
+    "Your model is ready at:\n",
+    "`/content/drive/MyDrive/stack-2.9/training_output/merged/`\n",
     "\n",
-    "**Note:** This model was trained on the full/mini dataset. For better performance, consider training on more data or more epochs.\n"
+    "Download it from Google Drive!"
    ]
   }
  ],
  "metadata": {
   "accelerator": "GPU",
   "colab": {
-   "name": "Stack 2.9 Colab Training (T4 GPU)",
+   "name": "Stack 2.9 Training",
    "provenance": []
   },
   "kernelspec": {