Fix colab training notebook for local Mac MPS training

- Fix JSON parsing error (extra closing braces)
- Convert from Colab/Google Drive to local training
- Update paths for local data
- Fix training and merge function calls
- Configure for MPS (Apple Silicon GPU)

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

colab_train_stack29.ipynb

@@ -4,26 +4,18 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "# 🚀 Stack 2.9 - Colab Training Notebook (Fixed)\n",
+    "# 🚀 Stack 2.9 - Local Training Notebook\n",
     "\n",
-    "**Zero-cost training on Google Colab free tier**\n",
+    "**Training on local Mac with MPS (Apple Silicon GPU)**\n",
     "\n",
-    "This notebook trains a LoRA adapter for Stack 2.9 Pattern Memory 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:** `./model_final/` (merged)\n",
+    "This notebook trains a LoRA adapter for Stack 2.9 on **Qwen2.5-Coder-7B** using your local machine.\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",
-    "\n",
-    "---"
+    "**Prerequisites:**\n",
+    "1. Base model downloaded to `./base_model_qwen7b/`\n",
+    "2. Run from the stack-2.9 directory\n",
+    "---\n"
    ]
   },
   {
@@ -32,46 +24,32 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Check GPU\n",
-    "!nvidia-smi"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## 1️⃣ Mount Google Drive (REQUIRED for persistence)\n",
-    "\n",
-    "Click the link, allow access, copy the auth code, paste it, and press Enter.\n",
+    "# Check system and GPU\n",
+    "import platform\n",
+    "import os\n",
     "\n",
-    "**Without Drive mounting, training will be lost if Colab disconnects!**"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from google.colab import drive\n",
-    "drive.mount('/content/drive')\n",
+    "print(\"=\"*60)\n",
+    "print(\"SYSTEM INFO\")\n",
+    "print(\"=\"*60)\n",
+    "print(f\"Platform: {platform.platform()}\")\n",
+    "print(f\"Python: {platform.python_version()}\")\n",
+    "print(f\"Working directory: {os.getcwd()}\")\n",
     "\n",
-    "# Set up paths on Drive - ALL OUTPUT GOES HERE\n",
-    "import os\n",
-    "BASE_PATH = \"/content/drive/MyDrive/stack-2.9\"\n",
-    "os.makedirs(BASE_PATH, exist_ok=True)\n",
-    "os.chdir(BASE_PATH)\n",
-    "print(f\"\\n✅ Working directory: {os.getcwd()}\")\n",
-    "print(f\"All outputs will be saved to: {BASE_PATH}\")\n",
-    "print(\"\\nCurrent folder contents:\")\n",
-    "!ls -la"
+    "try:\n",
+    "    import torch\n",
+    "    print(f\"PyTorch: {torch.__version__}\")\n",
+    "    print(f\"MPS available: {torch.backends.mps.is_available()}\")\n",
+    "    if torch.backends.mps.is_available():\n",
+    "        print(f\"MPS device: {torch.backends.mps.is_built()}\")\n",
+    "except ImportError:\n",
+    "    print(\"PyTorch not installed - run: pip install torch\")"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## 2️⃣ Clone Stack 2.9 Repository"
+    "## 1️⃣ Set Working Directory"
    ]
   },
   {
@@ -80,22 +58,35 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Clone into Drive if not already there\n",
-    "if not os.path.exists('stack-2.9'):\n",
-    "    !git clone https://github.com/my-ai-stack/stack-2.9.git\n",
+    "# Navigate to the stack-2.9 directory\n",
+    "import os\n",
     "\n",
-    "os.chdir('stack-2.9')\n",
-    "print(f\"Now in: {os.getcwd()}\")\n",
-    "!ls -la"
+    "STACK_PATH = \"/Users/walidsobhi/.openclaw/workspace/stack-2.9\"\n",
+    "os.chdir(STACK_PATH)\n",
+    "\n",
+    "BASE_MODEL_PATH = \"./base_model_qwen7b\"\n",
+    "DATA_PATH = \"./data/final/train.jsonl\"\n",
+    "OUTPUT_DIR = \"./training_output\"\n",
+    "\n",
+    "os.makedirs(OUTPUT_DIR, exist_ok=True)\n",
+    "\n",
+    "print(f\"✅ Working directory: {os.getcwd()}\")\n",
+    "print(f\"Base model path: {BASE_MODEL_PATH}\")\n",
+    "print(f\"Data path: {DATA_PATH}\")\n",
+    "print(f\"Output dir: {OUTPUT_DIR}\")\n",
+    "\n",
+    "# Check paths exist\n",
+    "print(\"\\n📁 Checking required files:\")\n",
+    "print(f\"  - Data exists: {os.path.exists(DATA_PATH)}\")\n",
+    "if not os.path.exists(DATA_PATH):\n",
+    "    print(f\"  ⚠️ Data not found at {DATA_PATH}\")"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## 3️⃣ Install Dependencies\n",
-    "\n",
-    "Takes 5-10 minutes."
+    "## 2️⃣ Install Dependencies"
    ]
   },
   {
@@ -104,9 +95,10 @@
    "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 bitsandbytes==0.43.0 datasets pyyaml\n",
+    "# Install required packages\n",
+    "!pip install torch torchvision torchaudio\n",
+    "!pip install transformers peft accelerate datasets pyyaml tqdm\n",
+    "!pip install scipy\n",
     "print(\"\\n✅ Dependencies installed\")"
    ]
   },
@@ -114,9 +106,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## 4️⃣ Create Mini Dataset (5K examples)\n",
-    "\n",
-    "Quick prototyping dataset - takes 1-2 minutes."
+    "## 3️⃣ Prepare Training Configuration"
    ]
   },
   {
@@ -125,19 +115,41 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "!python scripts/create_mini_dataset.py --size 5000 --output data_mini/train_mini.jsonl --source training-data/final/train.jsonl\n",
-    "\n",
-    "# Verify\n",
-    "!ls -lh data_mini/train_mini.jsonl"
+    "# Use local training config\n",
+    "import yaml\n",
+    "\n",
+    "config_path = \"./stack/training/train_config_local.yaml\"\n",
+    "\n",
+    "with open(config_path, 'r') as f:\n",
+    "    config = yaml.safe_load(f)\n",
+    "\n",
+    "# Update paths to match our setup\n",
+    "config['model']['name'] = BASE_MODEL_PATH\n",
+    "config['data']['input_path'] = DATA_PATH\n",
+    "config['output']['lora_dir'] = f\"{OUTPUT_DIR}/lora\"\n",
+    "config['output']['merged_dir'] = f\"{OUTPUT_DIR}/merged\"\n",
+    "config['hardware']['device'] = \"mps\"\n",
+    "\n",
+    "# Save updated config\n",
+    "updated_config_path = f\"{OUTPUT_DIR}/train_config.yaml\"\n",
+    "os.makedirs(OUTPUT_DIR, exist_ok=True)\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(\"\\nConfig summary:\")\n",
+    "print(f\"  - Model: {config['model']['name']}\")\n",
+    "print(f\"  - Data: {config['data']['input_path']}\")\n",
+    "print(f\"  - Device: {config['hardware']['device']}\")\n",
+    "print(f\"  - Epochs: {config['training']['num_epochs']}\")"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## 5️⃣ Prepare Training Configuration\n",
-    "\n",
-    "Use Colab-optimized config and point it to the mini dataset."
+    "## 4️⃣ Check Base Model"
    ]
   },
   {
@@ -146,31 +158,36 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Copy the Colab-optimized config\n",
-    "!cp stack_2_9_training/train_config_colab.yaml stack_2_9_training/train_config.yaml\n",
-    "\n",
-    "# Update config to use mini dataset paths\n",
-    "!sed -i 's|train_file: \"./training-data/train.jsonl\"|train_file: \"./data_mini/train_mini.jsonl\"|' stack_2_9_training/train_config.yaml\n",
-    "!sed -i 's|validation_file: \"./training-data/eval.jsonl\"|# validation_file: \"./data_mini/eval_mini.jsonl\"  # optional|' stack_2_9_training/train_config.yaml\n",
+    "# Check if base model exists\n",
+    "import os\n",
     "\n",
-    "print(\"✅ Config prepared for mini dataset\")\n",
-    "print(\"\\nTraining config (data section):\")\n",
-    "!grep -A3 'data:' stack_2_9_training/train_config.yaml"
+    "model_path = BASE_MODEL_PATH\n",
+    "print(f\"Checking base model at: {model_path}\")\n",
+    "\n",
+    "if os.path.exists(model_path):\n",
+    "    files = os.listdir(model_path)\n",
+    "    print(f\"✅ Base model found! {len(files)} files:\")\n",
+    "    for f in files[:10]:\n",
+    "        print(f\"   - {f}\")\n",
+    "    if len(files) > 10:\n",
+    "        print(f\"   ... and {len(files)-10} more\")\n",
+    "else:\n",
+    "    print(\"❌ Base model NOT found!\")\n",
+    "    print(\"\\nTo download Qwen2.5-Coder-7B:\")\n",
+    "    print(\"  huggingface-cli download Qwen/Qwen2.5-Coder-7B --local-dir ./base_model_qwen7b\")\n",
+    "    print(\"  OR\")\n",
+    "    print(\"  python -c \\\"from transformers import AutoModelForCausalLM; AutoModelForCausalLM.from_pretrained('Qwen/Qwen2.5-Coder-7B', local_dir='./base_model_qwen7b')\\\"\")"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## 6️⃣ Train LoRA Adapter\n",
+    "## 5️⃣ Train LoRA Adapter\n",
     "\n",
-    "⚠️ **This takes 3-5 hours. DO NOT INTERRUPT.**\n",
+    "⚠️ **This will take significant time on MPS.**\n",
     "\n",
-    "If Colab disconnects, reconnect and training will resume from checkpoint automatically.\n",
-    "\n",
-    "Watch for `Train loss:` decreasing. It should start ~2.0-3.0 and trend downward.\n",
-    "\n",
-    "Checkpoints saved every 500 steps to `./adapters_colab/` (on Drive)."
+    "MPS is slower than CUDA, so training will take longer. Consider using a smaller dataset for testing."
    ]
   },
   {
@@ -178,30 +195,13 @@
    "execution_count": null,
    "metadata": {},
    "outputs": [],
-   "source": [
-    "%env PYTHONUNBUFFERED=1\n",
-    "\n",
-    "print(\"\\n\" + \"=\"*60)\n",
-    "print(\"STARTING TRAINING\")\n",
-    "print(\"=\"*60)\n",
-    "print(f\"Working directory: {os.getcwd()}\")\n",
-    "print(f\"Config: stack_2_9_training/train_config.yaml\")\n",
-    "print(f\"Output will be saved to: ./adapters_colab/\")\n",
-    "print(\"=\"*60 + \"\\n\")\n",
-    "\n",
-    "!cd stack_2_9_training && python -m train_lora --config train_config.yaml\n",
-    "\n",
-    "print(\"\\n\" + \"=\"*60)\n",
-    "print(\"TRAINING FINISHED OR STOPPED\")\n",
-    "print(\"Check output above for errors or 'Training completed' message\")\n",
-    "print(\"=\"*60)"
-   ]
+   "source": "import os\nimport sys\n\n# Add the 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 using the correct function\nfrom train_lora import train_lora\n\ntrainer = train_lora(updated_config_path)\n\nprint(\"\\n\" + \"=\"*60)\nprint(\"TRAINING COMPLETED\")\nprint(\"=\"*60)"
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## 7️⃣ Verify Training Output"
+    "## 6️⃣ Verify Training Output"
    ]
   },
   {
@@ -210,17 +210,24 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "!ls -la adapters_colab/ 2>/dev/null || echo \"⚠️ adapters_colab/ not found - training may have failed or not run yet\"\n",
-    "!find adapters_colab -name \"*.bin\" -o -name \"*.safetensors\" 2>/dev/null | head -5 || echo \"No checkpoint files found\""
+    "lora_dir = f\"{OUTPUT_DIR}/lora\"\n",
+    "print(f\"Checking LoRA output: {lora_dir}\")\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\")"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## 8️⃣ Merge LoRA Adapter with Base Model\n",
-    "\n",
-    "Combines the trained adapter with the base model to produce a standalone fine-tuned model."
+    "## 7️⃣ Merge LoRA Adapter (Optional)"
    ]
   },
   {
@@ -228,18 +235,13 @@
    "execution_count": null,
    "metadata": {},
    "outputs": [],
-   "source": [
-    "!cd stack_2_9_training && python -m merge_adapter --base-model Qwen/Qwen2.5-Coder-7B\n",
-    "\n",
-    "print(\"\\n✅ Merged model created in ./model_final/\")\n",
-    "!ls -lh model_final/"
-   ]
+   "source": "# Merge LoRA with base model\nmerged_dir = f\"{OUTPUT_DIR}/merged\"\n\nprint(\"=\"*60)\nprint(\"MERGING LORA ADAPTER\")\nprint(\"=\"*60)\n\nsys.path.insert(0, './stack/training')\nfrom merge_adapter import merge_adapter\n\n# Create a minimal config for merging\nimport yaml\n\nmerge_config = {\n    'model': {'name': BASE_MODEL_PATH, 'trust_remote_code': True},\n    'output': {'lora_dir': f'{OUTPUT_DIR}/lora', 'merged_dir': merged_dir},\n    'quantization': {'enabled': False}\n}\n\n# Save config for merge\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}\")\nif os.path.exists(merged_dir):\n    print(\"Files:\", os.listdir(merged_dir)[:5])"
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## 9️⃣ Test Inference (Quick Check)"
+    "## 8️⃣ Test Inference"
    ]
   },
   {
@@ -251,20 +253,20 @@
     "from transformers import AutoTokenizer, AutoModelForCausalLM\n",
     "import torch\n",
     "\n",
-    "model_path = \"./model_final\"\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.bfloat16,\n",
-    "        device_map=\"auto\",\n",
+    "        torch_dtype=torch.float16,\n",
+    "        device_map=\"mps\",\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(model.device)\n",
+    "    inputs = tokenizer(prompt, return_tensors=\"pt\").to(\"mps\")\n",
     "    \n",
     "    print(\"Generating...\")\n",
     "    with torch.no_grad():\n",
@@ -277,13 +279,13 @@
     "        )\n",
     "    \n",
     "    response = tokenizer.decode(outputs[0], skip_special_tokens=True)\n",
-    "    print(\"=\\\"*40)\n",
+    "    print(\"=\"*40)\n",
     "    print(\"RESPONSE:\")\n",
-    "    print(\"=\\\"*40)\n",
+    "    print(\"=\"*40)\n",
     "    print(response[len(prompt):])\n",
     "except Exception as e:\n",
-    "    print(f\"❌ Error loading or running model: {e}\")\n",
-    "    print(\"\\nThis is expected if training hasn't completed yet or model files are missing.\")"
+    "    print(f\"❌ Error: {e}\")\n",
+    "    print(\"\\nThis is expected if training hasn't completed yet.\")"
    ]
   },
   {
@@ -292,25 +294,21 @@
    "source": [
     "## 🔚 Training Complete!\n",
     "\n",
-    "Your model is ready in `./model_final/` and saved to Google Drive.\n",
+    "Your model is ready in:\n",
+    "- LoRA adapter: `{OUTPUT_DIR}/lora/`\n",
+    "- Merged model: `{OUTPUT_DIR}/merged/`\n",
     "\n",
     "**Next steps:**\n",
-    "1. **Download** `model_final/` from Drive to your local machine\n",
-    "2. **Download datasets**: `python scripts/download_benchmark_datasets.py --benchmark both`\n",
-    "3. **Run evaluation**: `python stack-2.9-eval/run_proper_evaluation.py --benchmark humaneval --provider ollama --model ./model_final --k-samples 100`\n",
-    "4. **Upload model** to Hugging Face Hub\n",
-    "5. **Apply to Together AI**\n",
-    "\n",
-    "**Note:** This model was trained on 5K mini dataset (code completion only). For better tool capability, consider training on full dataset with synthetic tool data in a future run.\n",
-    "\n",
-    "**Need help?** See `COLAB_TRAINING.md` for detailed troubleshooting."
+    "1. Run evaluation on HumanEval/MBPP benchmarks\n",
+    "2. Upload to Hugging Face Hub\n",
+    "3. Apply to Together AI\n"
    ]
   }
  ],
  "metadata": {
   "accelerator": "GPU",
   "colab": {
-   "name": "Stack 2.9 Colab Training (Fixed)",
+   "name": "Stack 2.9 Local Training (Fixed)",
    "provenance": []
   },
   "kernelspec": {
@@ -323,5 +321,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 0
-}
-}
+}