fix: update Kaggle notebook with proper data handling and training configuration

- Fix data path issue: now correctly finds and uses training data
- Use repo's training-data/final/train.jsonl if available
- Fallback to create mini dataset (1K samples) from full data
- Removed broken model download step (assumes model in base_model_qwen7b)
- Proper config YAML generation with correct train_file path
- Better error handling and progress messages
- Simplified training and merge steps with proper imports

kaggle_train_stack29.ipynb

@@ -4,7 +4,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "# 🚀 Stack 2.9 - Kaggle Training Notebook\n",
+    "# 🚀 Stack 2.9 - Kaggle Training\n",
     "\n",
     "**Free GPU training on Kaggle**\n",
     "\n",
@@ -40,7 +40,31 @@
    "execution_count": null,
    "metadata": {},
    "outputs": [],
-   "source": "# STEP 2: Clone repo and setup paths\nimport os\nimport shutil\nimport subprocess\n\n# Change to a valid directory first (in case we're in a deleted folder)\nos.chdir(\"/kaggle/working\")\n\nREPO_DIR = \"/kaggle/working/stack-2.9\"\nMODEL_DIR = os.path.join(REPO_DIR, \"base_model_qwen7b\")\nOUTPUT_DIR = os.path.join(REPO_DIR, \"training_output\")\n\n# Remove old repo if exists (force fresh clone)\nif os.path.exists(REPO_DIR):\n    shutil.rmtree(REPO_DIR)\n\n# Clone fresh (now includes the input_path fix)\nsubprocess.run([\"git\", \"clone\", \"https://github.com/my-ai-stack/stack-2.9.git\", REPO_DIR], check=True)\nos.chdir(REPO_DIR)\n\nprint(f\"✅ Working in: {os.getcwd()}\")\nprint(f\"   MODEL_DIR: {MODEL_DIR}\")\nprint(f\"   OUTPUT_DIR: {OUTPUT_DIR}\")"
+   "source": [
+    "# STEP 2: Clone repo and setup paths\n",
+    "import os\n",
+    "import shutil\n",
+    "import subprocess\n",
+    "\n",
+    "# Change to a valid directory first (in case we're in a deleted folder)\n",
+    "os.chdir(\"/kaggle/working\")\n",
+    "\n",
+    "REPO_DIR = \"/kaggle/working/stack-2.9\"\n",
+    "MODEL_DIR = os.path.join(REPO_DIR, \"base_model_qwen7b\")\n",
+    "OUTPUT_DIR = os.path.join(REPO_DIR, \"training_output\")\n",
+    "\n",
+    "# Remove old repo if exists (force fresh clone)\n",
+    "if os.path.exists(REPO_DIR):\n",
+    "    shutil.rmtree(REPO_DIR)\n",
+    "\n",
+    "# Clone fresh (now includes the input_path fix)\n",
+    "subprocess.run([\"git\", \"clone\", \"https://github.com/my-ai-stack/stack-2.9.git\", REPO_DIR], check=True)\n",
+    "os.chdir(REPO_DIR)\n",
+    "\n",
+    "print(f\"✅ Working in: {os.getcwd()}\")\n",
+    "print(f\"   MODEL_DIR: {MODEL_DIR}\")\n",
+    "print(f\"   OUTPUT_DIR: {OUTPUT_DIR}\")"
+   ]
   },
   {
    "cell_type": "code",
@@ -62,26 +86,35 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# STEP 4: Download model (if not exists)\n",
-    "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
+    "# STEP 4: Prepare training data\n",
     "import os\n",
     "\n",
-    "if os.path.exists(os.path.join(MODEL_DIR, \"config.json\")):\n",
-    "    print(\"✅ Model already exists, skipping download!\")\n",
+    "# Check what training data is available\n",
+    "REPO_TRAIN_DATA = os.path.join(REPO_DIR, \"training-data/final/train.jsonl\")\n",
+    "MINI_DATA_DIR = os.path.join(REPO_DIR, \"data_mini\")\n",
+    "MINI_DATA_FILE = os.path.join(MINI_DATA_DIR, \"train_mini.jsonl\")\n",
+    "\n",
+    "print(\"🔍 Checking for training data...\")\n",
+    "if os.path.exists(REPO_TRAIN_DATA):\n",
+    "    print(f\"   Found full dataset: {REPO_TRAIN_DATA}\")\n",
+    "    # Create mini subset (1K samples) for faster training\n",
+    "    os.makedirs(MINI_DATA_DIR, exist_ok=True)\n",
+    "    if not os.path.exists(MINI_DATA_FILE):\n",
+    "        print(\"   Creating mini dataset (1000 samples)...\")\n",
+    "        import subprocess\n",
+    "        subprocess.run([\"python\", os.path.join(REPO_DIR, \"scripts/create_mini_dataset.py\"),\n",
+    "                       \"--size\", \"1000\", \"--output\", MINI_DATA_FILE, \"--source\", REPO_TRAIN_DATA], check=True)\n",
+    "    DATA_FILE = MINI_DATA_FILE\n",
     "else:\n",
-    "    print(\"⬇️ Downloading model (Qwen2.5-Coder-7B)...\")\n",
-    "    print(\"This takes ~10-15 minutes...\")\n",
-    "    \n",
-    "    tokenizer = AutoTokenizer.from_pretrained(\"Qwen/Qwen2.5-Coder-7B\", trust_remote_code=True)\n",
-    "    tokenizer.save_pretrained(MODEL_DIR)\n",
-    "    \n",
-    "    model = AutoModelForCausalLM.from_pretrained(\"Qwen/Qwen2.5-Coder-7B\", trust_remote_code=True)\n",
-    "    model.save_pretrained(MODEL_DIR)\n",
-    "    \n",
-    "    print(\"✅ Model downloaded!\")\n",
-    "\n",
-    "print(\"\\nModel files:\")\n",
-    "os.listdir(MODEL_DIR)"
+    "    print(\"   Full dataset not found, checking for existing mini dataset...\")\n",
+    "    if os.path.exists(MINI_DATA_FILE):\n",
+    "        DATA_FILE = MINI_DATA_FILE\n",
+    "        print(f\"   Using existing mini dataset: {MINI_DATA_FILE}\")\n",
+    "    else:\n",
+    "        raise FileNotFoundError(\"No training data found! Please ensure training-data/final/train.jsonl exists in the repo.\")\n",
+    "\n",
+    "print(f\"\\n✅ Using training data: {DATA_FILE}\")\n",
+    "print(f\"   Size: {os.path.getsize(DATA_FILE) / 1024:.1f} KB\")"
    ]
   },
   {
@@ -90,43 +123,72 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# STEP 5: Create config with train_dir and eval_dir\n",
+    "# STEP 5: Prepare config for training\n",
     "import yaml\n",
     "import os\n",
     "\n",
     "os.makedirs(OUTPUT_DIR, exist_ok=True)\n",
     "\n",
     "config = {\n",
-    "    'model': {'name': MODEL_DIR, 'trust_remote_code': True, 'torch_dtype': 'float16'},\n",
+    "    'model': {\n",
+    "        'name': 'Qwen/Qwen2.5-Coder-7B',\n",
+    "        'trust_remote_code': True,\n",
+    "        'torch_dtype': 'float16'\n",
+    "    },\n",
     "    'data': {\n",
-    "        'input_path': os.path.join(REPO_DIR, 'data/final/train.jsonl'),\n",
-    "        'train_dir': None,\n",
-    "        'eval_dir': None,\n",
+    "        'train_file': DATA_FILE,  # USE THE ACTUAL DATA FILE PATH\n",
     "        'max_length': 2048,\n",
-    "        'train_split': 0.9,\n",
-    "        'test_split': 0.1\n",
+    "        'train_split': 1.0  # Use all data for training\n",
+    "    },\n",
+    "    'lora': {\n",
+    "        'r': 16,\n",
+    "        'alpha': 32,\n",
+    "        'dropout': 0.05,\n",
+    "        'target_modules': ['q_proj', 'k_proj', 'v_proj', 'o_proj', 'gate_proj', 'up_proj', 'down_proj'],\n",
+    "        'bias': 'none',\n",
+    "        'task_type': 'CAUSAL_LM'\n",
+    "    },\n",
+    "    'training': {\n",
+    "        'num_epochs': 1,\n",
+    "        'batch_size': 2,\n",
+    "        'gradient_accumulation': 4,\n",
+    "        'learning_rate': 2e-4,\n",
+    "        'warmup_steps': 50,\n",
+    "        'weight_decay': 0.01,\n",
+    "        'max_grad_norm': 1.0,\n",
+    "        'logging_steps': 10,\n",
+    "        'save_steps': 100,\n",
+    "        'save_total_limit': 2,\n",
+    "        'fp16': True,\n",
+    "        'bf16': False,\n",
+    "        'gradient_checkpointing': True\n",
+    "    },\n",
+    "    'output': {\n",
+    "        'lora_dir': os.path.join(OUTPUT_DIR, 'lora'),\n",
+    "        'logging_dir': os.path.join(OUTPUT_DIR, 'logs')\n",
     "    },\n",
-    "    'lora': {'r': 16, 'alpha': 32, 'dropout': 0.05,\n",
-    "             'target_modules': ['q_proj', 'k_proj', 'v_proj', 'o_proj'],\n",
-    "             'bias': 'none', 'task_type': 'CAUSAL_LM'},\n",
-    "    'training': {'num_epochs': 1, 'batch_size': 2, 'gradient_accumulation': 4,\n",
-    "                 'learning_rate': 2e-4, 'warmup_steps': 50, 'weight_decay': 0.01,\n",
-    "                 'max_grad_norm': 1.0, 'logging_steps': 5, 'eval_steps': 100,\n",
-    "                 'save_steps': 200, 'save_total_limit': 2, 'fp16': True, 'bf16': False,\n",
-    "                 'gradient_checkpointing': True},\n",
-    "    'output': {'lora_dir': os.path.join(OUTPUT_DIR, 'lora'),\n",
-    "               'merged_dir': os.path.join(OUTPUT_DIR, 'merged')},\n",
-    "    'quantization': {'enabled': False},\n",
-    "    'hardware': {'device': 'cuda', 'num_gpus': 1, 'use_4bit': False, 'use_8bit': False}\n",
+    "    'quantization': {\n",
+    "        'enabled': False\n",
+    "    },\n",
+    "    'hardware': {\n",
+    "        'device': 'cuda',\n",
+    "        'num_gpus': 1,\n",
+    "        'use_4bit': False,\n",
+    "        'use_8bit': False\n",
+    "    }\n",
     "}\n",
     "\n",
     "config_path = os.path.join(OUTPUT_DIR, \"train_config.yaml\")\n",
     "with open(config_path, 'w') as f:\n",
-    "    yaml.dump(config, f)\n",
+    "    yaml.dump(config, f, default_flow_style=False)\n",
     "\n",
     "print(f\"✅ Config saved to: {config_path}\")\n",
-    "print(f\"   Device: {config['hardware']['device']}\")\n",
-    "print(f\"   Data: {config['data']['input_path']}\")"
+    "print(\"\\nConfig summary:\")\n",
+    "print(f\"   Model: {config['model']['name']}\")\n",
+    "print(f\"   Data: {config['data']['train_file']}\")\n",
+    "print(f\"   LoRA rank: {config['lora']['r']}\")\n",
+    "print(f\"   Batch size: {config['training']['batch_size']}\")\n",
+    "print(f\"   Epochs: {config['training']['num_epochs']}\")"
    ]
   },
   {
@@ -137,18 +199,28 @@
    "source": [
     "# STEP 6: Train LoRA\n",
     "import sys\n",
-    "sys.path.insert(0, os.path.join(REPO_DIR, \"stack/training\"))\n",
+    "sys.path.insert(0, os.path.join(REPO_DIR, \"stack_2_9_training\"))\n",
     "\n",
     "print(\"=\"*60)\n",
     "print(\"STARTING TRAINING\")\n",
     "print(\"=\"*60)\n",
+    "print(f\"Config: {config_path}\")\n",
+    "print(f\"Checkpoint dir: {config['output']['lora_dir']}\")\n",
+    "print(\"=\"*60 + \"\\n\")\n",
     "\n",
-    "from train_lora import train_lora\n",
-    "trainer = train_lora(config_path)\n",
+    "# Import and run training\n",
+    "from stack_2_9_training.train_lora import train_lora\n",
     "\n",
-    "print(\"=\"*60)\n",
-    "print(\"TRAINING COMPLETED!\")\n",
-    "print(\"=\"*60)"
+    "try:\n",
+    "    trainer = train_lora(config_path)\n",
+    "    print(\"\\n\" + \"=\"*60)\n",
+    "    print(\"TRAINING COMPLETED SUCCESSFULLY\")\n",
+    "    print(\"=\"*60)\n",
+    "except Exception as e:\n",
+    "    print(f\"\\n❌ Training failed: {e}\")\n",
+    "    import traceback\n",
+    "    traceback.print_exc()\n",
+    "    raise"
    ]
   },
   {
@@ -157,43 +229,41 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# STEP 7: Merge model\n",
+    "# STEP 7: Merge LoRA adapter with base model\n",
     "import sys\n",
-    "sys.path.insert(0, os.path.join(REPO_DIR, \"stack/training\"))\n",
-    "from merge_adapter import merge_adapter\n",
+    "sys.path.insert(0, os.path.join(REPO_DIR, \"stack_2_9_training\"))\n",
+    "from stack_2_9_training.merge_adapter import merge_adapter\n",
     "\n",
-    "merged_dir = os.path.join(OUTPUT_DIR, \"merged\")\n",
+    "lora_dir = config['output']['lora_dir']\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': os.path.join(OUTPUT_DIR, 'lora'), 'merged_dir': merged_dir},\n",
-    "    'quantization': {'enabled': False}\n",
-    "}\n",
-    "\n",
-    "merge_cfg_path = os.path.join(OUTPUT_DIR, \"merge_config.yaml\")\n",
-    "with open(merge_cfg_path, 'w') as f:\n",
-    "    yaml.dump(merge_config, f)\n",
+    "print(\"=\"*60)\n",
+    "print(\"MERGING LORA ADAPTER\")\n",
+    "print(\"=\"*60)\n",
+    "print(f\"LoRA adapter: {lora_dir}\")\n",
+    "print(f\"Output: {merged_dir}\")\n",
     "\n",
-    "merge_adapter(merge_cfg_path, os.path.join(OUTPUT_DIR, \"lora\"), merged_dir)\n",
+    "try:\n",
+    "    merge_adapter(\n",
+    "        base_model_name_or_path=config['model']['name'],\n",
+    "        adapter_path=lora_dir,\n",
+    "        output_path=merged_dir,\n",
+    "        use_safetensors=True\n",
+    "    )\n",
+    "    print(\"\\n✅ Merge completed!\")\n",
+    "    print(f\"Merged model files: {os.listdir(merged_dir)}\")\n",
+    "except Exception as e:\n",
+    "    print(f\"\\n❌ Merge failed: {e}\")\n",
+    "    import traceback\n",
+    "    traceback.print_exc()\n",
+    "    raise\n",
     "\n",
-    "print(f\"✅ Merged model saved to: {merged_dir}\")\n",
-    "print(\"Files:\", os.listdir(merged_dir))"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# STEP 8: Done!\n",
     "print(\"=\"*60)\n",
-    "print(\"🎉 TRAINING COMPLETE!\")\n",
+    "print(\"🎉 ALL DONE!\")\n",
     "print(\"=\"*60)\n",
-    "print(f\"LoRA adapter: {os.path.join(OUTPUT_DIR, 'lora')}\")\n",
-    "print(f\"Merged model: {os.path.join(OUTPUT_DIR, 'merged')}\")\n",
-    "print(\"\\n📥 Download from: Kaggle → Output tab\")"
+    "print(f\"\\n📦 Merged model ready at: {merged_dir}\")\n",
+    "print(\"\\n⏳ Download the 'merged' folder from Kaggle's Output panel before the session ends!\")"
    ]
   }
  ],
@@ -210,4 +280,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 0
-}
+}