{
 "nbformat": 4,
 "nbformat_minor": 5,
 "metadata": {
  "colab": {
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  },
  "accelerator": "GPU"
 },
 "cells": [
  {
   "cell_type": "markdown",
   "id": "hdr",
   "metadata": {},
   "source": [
    "#  Data-Centric AI Agent — Training Notebook\n",
    "\n",
    "**OpenEnv Hackathon 2026**\n",
    "\n",
    "Trains **Qwen2.5-1.5B-Instruct** (4-bit QLoRA) to orchestrate data-cleaning specialists\n",
    "via GRPO reinforcement learning on the Data-Centric AI environment.\n",
    "\n",
    "| | |\n",
    "|---|---|\n",
    "| **Space** | https://huggingface.co/spaces/Aswini-Kumar/data-centric-env |\n",
    "| **Repo** | https://github.com/CelestialWorthyOfHeavenAndEarth/data-centric-env |\n",
    "| **Algorithm** | SFT warmup → GRPO (TRL) |\n",
    "| **Model** | Qwen/Qwen2.5-1.5B-Instruct, 4-bit QLoRA r=8 |\n",
    "\n",
    "**Pipeline:**\n",
    "1. Install deps\n",
    "2. Clone repo + start env server\n",
    "3. Generate SFT warmup data\n",
    "4. Phase 1 — SFT warmup (~15 min)\n",
    "5. Phase 2 — GRPO training (~45-90 min depending on hardware)\n",
    "6. Plot reward curves\n",
    "7. Evaluate vs baselines\n",
    "8. Save model\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "s1_md",
   "metadata": {},
   "source": [
    "##  Step 1 — Install Dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "s1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Pin transformers first to avoid torchao conflict\n",
    "# (transformers ≥ 4.52 requires torchao ≥ 0.9 which needs torch.int1 — not in Colab yet)\n",
    "!pip install \"transformers>=4.44.0,<4.52.0\" --quiet\n",
    "\n",
    "# Unsloth — detects the installed torch version automatically\n",
    "!pip install \"unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git\" --quiet\n",
    "\n",
    "# Training stack\n",
    "!pip install trl>=0.15.0 datasets>=2.0.0 accelerate>=0.30.0 --quiet\n",
    "!pip install scikit-learn>=1.3.0 pandas>=2.0.0 numpy matplotlib --quiet\n",
    "!pip install \"openenv-core[core]>=0.2.1\" --quiet\n",
    "\n",
    "# Verify Unsloth\n",
    "from unsloth import FastLanguageModel\n",
    "import torch\n",
    "print(f' Unsloth ready | torch {torch.__version__} | GPU: {torch.cuda.get_device_name(0) if torch.cuda.is_available() else \"CPU\"}')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "s2_md",
   "metadata": {},
   "source": [
    "##  Step 2 — Clone Repo & Start Environment Server"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "s2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os, sys, shutil, subprocess, time, requests\n",
    "\n",
    "REPO     = 'https://github.com/CelestialWorthyOfHeavenAndEarth/data-centric-env.git'\n",
    "WORK_DIR = '/content/data-centric-env'\n",
    "\n",
    "# Clone or update\n",
    "if os.path.exists(f'{WORK_DIR}/pyproject.toml'):\n",
    "    print('Repo exists — pulling latest...')\n",
    "    os.system(f'git -C {WORK_DIR} pull origin main')\n",
    "else:\n",
    "    if os.path.exists(WORK_DIR): shutil.rmtree(WORK_DIR)\n",
    "    os.system(f'git clone {REPO} {WORK_DIR}')\n",
    "\n",
    "os.chdir(WORK_DIR)\n",
    "sys.path.insert(0, WORK_DIR)\n",
    "os.makedirs('logs', exist_ok=True)\n",
    "os.makedirs('plots', exist_ok=True)\n",
    "print('CWD:', os.getcwd())\n",
    "os.system('git log --oneline -3')\n",
    "os.system('pip install -e . --quiet 2>/dev/null || true')\n",
    "print(' Repo ready')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "s2b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import subprocess, time, requests, os\n",
    "\n",
    "# Kill any leftover server\n",
    "os.system('fuser -k 8000/tcp 2>/dev/null || true')\n",
    "time.sleep(2)\n",
    "\n",
    "server = subprocess.Popen(\n",
    "    ['python', '-m', 'uvicorn', 'server.app:app',\n",
    "     '--host', '0.0.0.0', '--port', '8000', '--log-level', 'warning'],\n",
    "    stdout=subprocess.PIPE, stderr=subprocess.PIPE\n",
    ")\n",
    "\n",
    "for i in range(40):\n",
    "    try:\n",
    "        r = requests.get('http://localhost:8000/health', timeout=2)\n",
    "        if r.status_code == 200:\n",
    "            print(f' Server ready after {i+1}s — {r.json()}')\n",
    "            break\n",
    "    except: time.sleep(1)\n",
    "else:\n",
    "    out, err = server.communicate(timeout=5)\n",
    "    print('STDOUT:', out.decode()[-2000:])\n",
    "    print('STDERR:', err.decode()[-2000:])\n",
    "    raise RuntimeError('Server failed to start')\n",
    "\n",
    "os.environ['ENV_URL'] = 'http://localhost:8000'\n",
    "print('ENV_URL set to http://localhost:8000')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "s3_md",
   "metadata": {},
   "source": [
    "##  Step 3 — Generate SFT Warmup Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "s3",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os, time\n",
    "os.chdir('/content/data-centric-env')\n",
    "\n",
    "sft_path = 'sft_data.jsonl'\n",
    "if os.path.exists(sft_path):\n",
    "    count = sum(1 for _ in open(sft_path))\n",
    "    age   = (time.time() - os.path.getmtime(sft_path)) / 60\n",
    "    print(f' sft_data.jsonl: {count} examples ({age:.0f} min old) — reusing')\n",
    "else:\n",
    "    print('Generating SFT warmup data (~2 min)...')\n",
    "    os.system('python sft_generator.py')\n",
    "    count = sum(1 for _ in open(sft_path))\n",
    "    print(f' Generated {count} SFT examples')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "s4_md",
   "metadata": {},
   "source": [
    "##  Step 4 — Load Qwen2.5-1.5B (4-bit QLoRA)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "s4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os, sys, torch\n",
    "os.chdir('/content/data-centric-env')\n",
    "sys.path.insert(0, '/content/data-centric-env')\n",
    "\n",
    "from unsloth import FastLanguageModel\n",
    "from train_data_centric import load_model\n",
    "\n",
    "model, tokenizer = load_model()\n",
    "\n",
    "vram  = torch.cuda.memory_allocated() / 1e9\n",
    "total = torch.cuda.get_device_properties(0).total_memory / 1e9\n",
    "print(f'\\n VRAM: {vram:.1f} GB used / {total:.1f} GB total ({100*vram/total:.0f}%)')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "s5_md",
   "metadata": {},
   "source": [
    "##  Step 5 — Phase 1: SFT Warmup\n",
    "\n",
    "One epoch of supervised fine-tuning on ~9,480 heuristic trajectories.\n",
    "Teaches the model valid command syntax before RL starts.\n",
    "Without this, the model outputs gibberish and gets zero reward.\n",
    "\n",
    "**TensorBoard:** Run Step 6 to open the live dashboard.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "s5",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os, sys\n",
    "os.chdir('/content/data-centric-env')\n",
    "sys.path.insert(0, '/content/data-centric-env')\n",
    "\n",
    "from train_data_centric import run_sft_warmup\n",
    "\n",
    "model = run_sft_warmup(model, tokenizer)\n",
    "print(' SFT warmup complete — model knows command syntax')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "s6_md",
   "metadata": {},
   "source": [
    "##  Step 6 — TensorBoard (Experiment Tracking)\n",
    "\n",
    "Launch this now — it will show SFT logs immediately and GRPO logs as training progresses.\n",
    "Refresh the TensorBoard tab every few minutes during Step 7.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "s6",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext tensorboard\n",
    "%tensorboard --logdir /content/data-centric-env/logs\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "s7_md",
   "metadata": {},
   "source": [
    "##  Step 7 — Phase 2: GRPO Training\n",
    "\n",
    "Agent learns *when* to use each command via reinforcement learning.\n",
    "Each GRPO step: model generates a command → environment executes it → reward flows back.\n",
    "\n",
    "**Reward discrimination:** The reward is now strictly discriminating:\n",
    "- Hitting target efficiently (few steps): **+0.7 → +0.9**\n",
    "- Hitting target slowly (many steps): **+0.3 → +0.5**\n",
    "- Failing to hit target: **−0.1 → −0.5**\n",
    "- Blind apply / blind submit: **large penalties**\n",
    "\n",
    "Checkpoints saved every 25 steps to `./data-centric-checkpoints/`.\n",
    "If the runtime disconnects, re-run all cells — it auto-resumes from the latest checkpoint.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "s7",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os, sys, glob\n",
    "os.chdir('/content/data-centric-env')\n",
    "sys.path.insert(0, '/content/data-centric-env')\n",
    "\n",
    "from train_data_centric import run_grpo_training\n",
    "\n",
    "# Auto-resume from latest checkpoint\n",
    "ckpt_dir = './data-centric-checkpoints'\n",
    "resume_from = None\n",
    "if os.path.exists(ckpt_dir):\n",
    "    checkpoints = sorted(glob.glob(f'{ckpt_dir}/checkpoint-*'))\n",
    "    if checkpoints:\n",
    "        resume_from = checkpoints[-1]\n",
    "        print(f'Resuming from: {resume_from}')\n",
    "\n",
    "model = run_grpo_training(model, tokenizer, resume_from_checkpoint=resume_from)\n",
    "print(' GRPO training complete')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "s8_md",
   "metadata": {},
   "source": [
    "##  Step 8 — Plot Reward Curves\n",
    "\n",
    "Generates 4 plots from the training log:\n",
    "- **reward_curve.png** — Episode reward over time with rolling mean\n",
    "- **success_rate.png** — Success rate per curriculum level\n",
    "- **accuracy_gain.png** — Accuracy gain vs baseline per episode\n",
    "- **curriculum.png** — Curriculum level progression\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "s8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.chdir('/content/data-centric-env')\n",
    "\n",
    "log_path = 'logs/training.jsonl'\n",
    "\n",
    "if not os.path.exists(log_path):\n",
    "    print('No training log yet — run Step 7 first')\n",
    "else:\n",
    "    n_eps = sum(1 for _ in open(log_path))\n",
    "    print(f'Training log: {n_eps} episodes')\n",
    "\n",
    "    os.system(f'python plot_rewards.py --log {log_path} --out plots/')\n",
    "\n",
    "    from IPython.display import Image, display\n",
    "    for fname in ['reward_curve.png', 'success_rate.png', 'accuracy_gain.png', 'curriculum.png']:\n",
    "        path = f'plots/{fname}'\n",
    "        if os.path.exists(path):\n",
    "            print(f'\\n--- {fname} ---')\n",
    "            display(Image(path, width=900))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "s9_md",
   "metadata": {},
   "source": [
    "##  Step 9 — Evaluate: Trained Agent vs Baselines\n",
    "\n",
    "Runs the trained agent against:\n",
    "- **Random baseline** — random command each step\n",
    "- **Heuristic baseline** — fixed rule-based sequence\n",
    "- **Trained agent** — your GRPO-trained model\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "s9",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os, json\n",
    "os.chdir('/content/data-centric-env')\n",
    "\n",
    "os.system('python eval_data_centric.py')\n",
    "\n",
    "if os.path.exists('eval_results.json'):\n",
    "    with open('eval_results.json') as f:\n",
    "        results = json.load(f)\n",
    "    print('\\n=== Evaluation Results ===')\n",
    "    for k, v in results.items():\n",
    "        print(f'  {k}: {v}')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "s10_md",
   "metadata": {},
   "source": [
    "##  Step 10 — Save Trained LoRA Adapter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "s10",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os, sys\n",
    "os.chdir('/content/data-centric-env')\n",
    "sys.path.insert(0, '/content/data-centric-env')\n",
    "\n",
    "from train_data_centric import save_model\n",
    "save_model(model, tokenizer)\n",
    "print(' Adapter saved to ./data-centric-adapter/')\n",
    "\n",
    "# List outputs\n",
    "import glob\n",
    "print('\\n Training outputs:')\n",
    "for f in sorted(glob.glob('plots/*.png') + glob.glob('logs/*.jsonl') + ['eval_results.json']):\n",
    "    if os.path.exists(f):\n",
    "        size = os.path.getsize(f) / 1024\n",
    "        print(f'  {f} ({size:.1f} KB)')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "s11_md",
   "metadata": {},
   "source": [
    "##  Step 11 — Download Results\n",
    "\n",
    "Download plots and training log to your local machine.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "s11",
   "metadata": {},
   "outputs": [],
   "source": [
    "from google.colab import files\n",
    "import glob, os\n",
    "\n",
    "to_download = glob.glob('plots/*.png') + ['logs/training.jsonl', 'eval_results.json']\n",
    "\n",
    "for f in to_download:\n",
    "    if os.path.exists(f):\n",
    "        print(f'Downloading: {f}')\n",
    "        files.download(f)\n"
   ]
  }
 ]
}