remove: old REINFORCE notebook, keeping only real GRPO training notebook

GRPO_Training_Pipeline.ipynb

@@ -1,98 +0,0 @@
-{
- "metadata": {
-  "kernelspec": {
-   "language": "python",
-   "display_name": "Python 3",
-   "name": "python3"
-  },
-  "language_info": {
-   "name": "python",
-   "version": "3.12.12",
-   "mimetype": "text/x-python",
-   "codemirror_mode": {
-    "name": "ipython",
-    "version": 3
-   },
-   "pygments_lexer": "ipython3",
-   "nbconvert_exporter": "python",
-   "file_extension": ".py"
-  },
-  "kaggle": {
-   "accelerator": "nvidiaTeslaT4",
-   "dataSources": [],
-   "dockerImageVersionId": 31329,
-   "isInternetEnabled": true,
-   "language": "python",
-   "sourceType": "notebook",
-   "isGpuEnabled": true
-  }
- },
- "nbformat_minor": 4,
- "nbformat": 4,
- "cells": [
-  {
-   "cell_type": "code",
-   "source": "!pip install \"unsloth[kaggle-new] @ git+https://github.com/unslothai/unsloth.git\" -q\n!pip install trl datasets -q",
-   "metadata": {
-    "trusted": true,
-    "execution": {
-     "iopub.status.busy": "2026-04-25T04:07:23.495166Z",
-     "iopub.execute_input": "2026-04-25T04:07:23.495514Z",
-     "iopub.status.idle": "2026-04-25T04:08:11.862767Z",
-     "shell.execute_reply.started": "2026-04-25T04:07:23.495482Z",
-     "shell.execute_reply": "2026-04-25T04:08:11.861936Z"
-    }
-   },
-   "outputs": [],
-   "execution_count": null
-  },
-  {
-   "cell_type": "code",
-   "source": "import os, re, json, time, requests, random, gc\nimport torch\nimport matplotlib.pyplot as plt\nfrom unsloth import FastLanguageModel\n\n# HF Login (set HF_TOKEN in Kaggle Secrets → Add-ons → Secrets)\ntry:\n    from kaggle_secrets import UserSecretsClient\n    hf_token = UserSecretsClient().get_secret(\"HF_TOKEN\")\n    from huggingface_hub import login\n    login(token=hf_token)\n    print(\"✅ Logged into HuggingFace\")\nexcept:\n    print(\"⚠️ HF_TOKEN not found — model push will fail later\")\n\nENV_URL = \"https://cobedigger-dataengenv.hf.space\"\n\nprint(\"Waking up HF Space...\")\nfor attempt in range(15):\n    try:\n        r = requests.get(f\"{ENV_URL}/health\", timeout=60)\n        if r.status_code == 200:\n            print(f\"✅ Space is live: {r.json()}\")\n            break\n    except Exception as e:\n        print(f\"  Attempt {attempt+1}: {e}\")\n    time.sleep(10)\nelse:\n    print(\"❌ Space did not wake up!\")\n",
-   "metadata": {
-    "trusted": true,
-    "execution": {
-     "iopub.status.busy": "2026-04-25T04:08:29.774532Z",
-     "iopub.execute_input": "2026-04-25T04:08:29.775094Z",
-     "iopub.status.idle": "2026-04-25T04:09:04.127730Z",
-     "shell.execute_reply.started": "2026-04-25T04:08:29.775053Z",
-     "shell.execute_reply": "2026-04-25T04:09:04.126832Z"
-    }
-   },
-   "outputs": [],
-   "execution_count": null
-  },
-  {
-   "cell_type": "code",
-   "source": "model, tokenizer = FastLanguageModel.from_pretrained(\n    model_name=\"unsloth/Meta-Llama-3.1-8B-Instruct-bnb-4bit\",\n    max_seq_length=2048,\n    dtype=None,\n    load_in_4bit=True,\n)\nmodel = FastLanguageModel.get_peft_model(\n    model, r=16,\n    target_modules=[\"q_proj\",\"k_proj\",\"v_proj\",\"o_proj\",\"gate_proj\",\"up_proj\",\"down_proj\"],\n    lora_alpha=16, lora_dropout=0, bias=\"none\",\n    use_gradient_checkpointing=\"unsloth\", random_state=42,\n)\nprint(f\"✅ Model loaded on {torch.cuda.get_device_name(0)}\")\nprint(f\"   VRAM used: {torch.cuda.memory_allocated()/1024**3:.1f} GB\")\n",
-   "metadata": {
-    "trusted": true,
-    "execution": {
-     "iopub.status.busy": "2026-04-25T04:09:40.052566Z",
-     "iopub.execute_input": "2026-04-25T04:09:40.053377Z",
-     "iopub.status.idle": "2026-04-25T04:10:18.432112Z",
-     "shell.execute_reply.started": "2026-04-25T04:09:40.053331Z",
-     "shell.execute_reply": "2026-04-25T04:10:18.431319Z"
-    }
-   },
-   "outputs": [],
-   "execution_count": null
-  },
-  {
-   "cell_type": "code",
-   "source": "###############################################################\n# REAL GRPO TRAINING\n###############################################################\n\nimport torch, requests, json, time, gc, re\nimport torch.nn.functional as F\nimport matplotlib\nmatplotlib.use(\"Agg\")\nimport matplotlib.pyplot as plt\nfrom unsloth import FastLanguageModel\n\nENV = \"https://cobedigger-dataengenv.hf.space\"\n\nG           = 4\nTRAIN_STEPS = 25\nLR          = 1e-5   # lower LR — 4bit models are sensitive\nMAX_NEW     = 150\nMAX_PROMPT  = 350\nMAX_FULL    = MAX_PROMPT + MAX_NEW\n\nSYSTEM = (\n    \"You are an ML engineer. Respond with ONLY a JSON object, nothing else.\\n\"\n    \"Format: {\\\"action_type\\\": \\\"<action>\\\", \\\"payload\\\": {}}\\n\"\n    \"Actions: inspect_data, run_script, edit_script, query_actor, submit\\n\"\n    \"For edit_script: {\\\"action_type\\\": \\\"edit_script\\\", \\\"payload\\\": {\\\"old\\\": \\\"...\\\", \\\"new\\\": \\\"...\\\"}}\\n\\n\"\n    \"Stage 1 bug: column 'age_years' should be 'age', missing df.dropna()\\n\"\n    \"Stage 2 bug: MLPClassifier needs StandardScaler\\n\"\n    \"Stage 3 bug: scaler.fit called before train_test_split\\n\"\n    \"Stage 4 bug: LogisticRegression needs class_weight='balanced'\"\n)\n\nSTAGE_PROMPTS = {\n    1: \"Stage 1 error: KeyError 'age_years'. Column should be 'age'. NaN rows present. Pick one action.\",\n    2: \"Stage 2 error: MLPClassifier loss is NaN. Needs StandardScaler normalization. Pick one action.\",\n    3: \"Stage 3: accuracy 0.98 — scaler fit before split causes leakage. Pick one action.\",\n    4: \"Stage 4: fairness_gap too high — needs class_weight='balanced'. Pick one action.\",\n}\n\nSTAGE_HINTS = {\n    1: [\"age_years\", \"age\", \"dropna\", \"rename\"],\n    2: [\"StandardScaler\", \"scaler\", \"normalize\"],\n    3: [\"train_test_split\", \"leakage\", \"after\"],\n    4: [\"class_weight\", \"balanced\"],\n}\n\nVALID_ACTIONS = {\"inspect_data\", \"run_script\", \"edit_script\", \"query_actor\", \"submit\"}\n\ndef parse_action(text):\n    s = text.find(\"{\")\n    if s != -1:\n        depth = 0\n        for i, ch in enumerate(text[s:], s):\n            if ch == \"{\": depth += 1\n            elif ch == \"}\":\n                depth -= 1\n                if depth == 0:\n                    try:\n                        return json.loads(text[s:i+1])\n                    except Exception:\n                        break\n    m = re.search(r'\"?action_?type\"?\\s*[=:]\\s*\"?(\\w+)\"?', text, re.IGNORECASE)\n    if m:\n        atype = m.group(1).lower().replace(\" \", \"_\")\n        for va in VALID_ACTIONS:\n            if va.replace(\"_\",\"\") in atype.replace(\"_\",\"\"):\n                return {\"action_type\": va, \"payload\": {}}\n    return None\n\ndef score_action(text, stage):\n    action = parse_action(text)\n    if action is None:\n        return 0.0\n    action.setdefault(\"payload\", {})\n    if isinstance(action[\"payload\"], str):\n        action[\"payload\"] = {\"old\": action[\"payload\"], \"new\": \"\"}\n\n    atype = str(action.get(\"action_type\",\"\")).lower().replace(\" \",\"_\")\n    for va in VALID_ACTIONS:\n        if va.replace(\"_\",\"\") in atype.replace(\"_\",\"\"):\n            atype = va; break\n    else:\n        return 0.05\n    action[\"action_type\"] = atype\n\n    reward = 0.15\n    good = {1:{\"inspect_data\",\"edit_script\",\"run_script\"},\n            2:{\"inspect_data\",\"edit_script\",\"run_script\"},\n            3:{\"run_script\",\"edit_script\"},\n            4:{\"query_actor\",\"edit_script\",\"run_script\"}}\n    if atype in good.get(stage, set()):\n        reward = 0.25\n\n    if atype == \"edit_script\":\n        full = (text + json.dumps(action[\"payload\"])).lower()\n        hits = sum(1 for h in STAGE_HINTS.get(stage,[]) if h.lower() in full)\n        if hits >= 1: reward = 0.40\n        if hits >= 2: reward = 0.55\n        try:\n            requests.post(f\"{ENV}/reset\", timeout=15)\n            r = requests.post(f\"{ENV}/step\", json=action, timeout=15).json()\n            sc = float(r.get(\"reward\",{}).get(\"score\", 0.0))\n            if sc > 0: reward = min(1.0, 0.55 + sc * 0.45)\n        except Exception:\n            pass\n    return reward\n\ndef make_prompt(stage):\n    msgs = [{\"role\":\"system\",\"content\":SYSTEM},\n            {\"role\":\"user\",\"content\":STAGE_PROMPTS[stage]}]\n    return tokenizer.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True)\n\ndef reinforce_loss(input_ids, attention_mask, advantages):\n    logits  = model(input_ids=input_ids, attention_mask=attention_mask).logits\n    shift_l = logits[..., :-1, :].contiguous()\n    shift_b = input_ids[..., 1:].contiguous()\n    lp      = F.log_softmax(shift_l, dim=-1)\n    tok_lp  = lp.gather(-1, shift_b.unsqueeze(-1)).squeeze(-1)\n    mask    = attention_mask[..., 1:].float()\n    seq_lp  = (tok_lp * mask).sum(-1) / mask.sum(-1).clamp(min=1)\n    return -(advantages * seq_lp).mean()\n\n# ── Setup — stay in training mode the whole time ─────────────\nprint(\"Setting up GRPO training...\")\ntorch.cuda.empty_cache(); gc.collect()\nFastLanguageModel.for_training(model)\nmodel.train()\n\noptimizer = torch.optim.AdamW(\n    [p for p in model.parameters() if p.requires_grad], lr=LR\n)\n\ngrpo_rewards, grpo_losses = [], []\nprint(f\"\\nStarting GRPO: {TRAIN_STEPS} steps, G={G} completions/step\")\nprint(\"=\" * 60)\n\nfor step in range(TRAIN_STEPS):\n    stage  = [1, 2, 3, 4][step % 4]\n    prompt = make_prompt(stage)\n\n    enc = tokenizer(prompt, return_tensors=\"pt\",\n                    truncation=True, max_length=MAX_PROMPT).to(\"cuda\")\n    prompt_len = enc[\"input_ids\"].shape[1]\n\n    # 1. Generate with torch.no_grad() — stay in train mode, just no grad tracking\n    with torch.no_grad():\n        outs = model.generate(\n            **enc,\n            max_new_tokens=MAX_NEW,\n            do_sample=True,\n            temperature=0.9,\n            num_return_sequences=G,\n            pad_token_id=tokenizer.eos_token_id,\n            repetition_penalty=1.1,\n        )\n    completions = [\n        tokenizer.decode(outs[i][prompt_len:], skip_special_tokens=True)\n        for i in range(G)\n    ]\n\n    # 2. Score\n    rewards = [score_action(c, stage) for c in completions]\n    time.sleep(0.2)\n\n    rewards_t = torch.tensor(rewards, dtype=torch.float32)\n    mean_r    = rewards_t.mean()\n    std_r     = rewards_t.std().clamp(min=1e-6)\n\n    print(f\"  Step {step+1:02d}/{TRAIN_STEPS} | Stage {stage} | \"\n          f\"Rewards={[f'{r:.2f}' for r in rewards]} | Avg={mean_r.item():.3f}\", end=\"\")\n\n    if std_r < 1e-4:\n        print(\" | SKIP (no variance)\")\n        grpo_rewards.append(mean_r.item())\n        grpo_losses.append(0.0)\n        continue\n\n    advantages = ((rewards_t - mean_r) / std_r).to(\"cuda\")\n\n    # 3. Tokenise\n    full_ids_list, full_mask_list = [], []\n    for c in completions:\n        ef = tokenizer(prompt + c, return_tensors=\"pt\",\n                       truncation=True, padding=\"max_length\", max_length=MAX_FULL)\n        full_ids_list.append(ef[\"input_ids\"])\n        full_mask_list.append(ef[\"attention_mask\"])\n    full_ids  = torch.cat(full_ids_list,  dim=0).to(\"cuda\")\n    full_mask = torch.cat(full_mask_list, dim=0).to(\"cuda\")\n\n    # 4. Backprop — model already in train mode, no mode switch needed\n    loss = reinforce_loss(full_ids, full_mask, advantages)\n    loss.backward()\n    torch.nn.utils.clip_grad_norm_(\n        [p for p in model.parameters() if p.requires_grad], 0.5\n    )\n    optimizer.step()\n    optimizer.zero_grad()\n    torch.cuda.empty_cache()\n\n    grpo_rewards.append(mean_r.item())\n    grpo_losses.append(loss.item())\n    print(f\" | Loss={loss.item():.4f}\")\n\nprint(\"\\nGRPO training complete!\")\nprint(f\"   First 5 avg reward : {sum(grpo_rewards[:5])/5:.3f}\")\nprint(f\"   Last  5 avg reward : {sum(grpo_rewards[-5:])/5:.3f}\")\n\nfig, (ax1, ax2) = plt.subplots(1, 2, figsize=(13, 4))\nax1.plot(grpo_rewards, color=\"#22c55e\", linewidth=2, marker=\"o\", markersize=4)\nax1.set_title(\"GRPO Avg Reward per Step\", fontweight=\"bold\")\nax1.set_xlabel(\"Training Step\"); ax1.set_ylabel(\"Avg Group Reward\")\nax1.set_ylim(-0.05, 1.1); ax1.grid(True, alpha=0.3)\nax2.plot(grpo_losses, color=\"#f59e0b\", linewidth=2, marker=\"o\", markersize=4)\nax2.set_title(\"GRPO Policy Loss per Step\", fontweight=\"bold\")\nax2.set_xlabel(\"Training Step\"); ax2.set_ylabel(\"Loss\")\nax2.grid(True, alpha=0.3)\nplt.tight_layout()\nplt.savefig(\"grpo_training_curve.png\", dpi=150, bbox_inches=\"tight\")\nplt.show()\n\nwith open(\"grpo_rewards.json\", \"w\") as f:\n    json.dump({\"rewards\": grpo_rewards, \"losses\": grpo_losses}, f)\nprint(\"Saved grpo_training_curve.png and grpo_rewards.json\")\n",
-   "metadata": {
-    "trusted": true
-   },
-   "outputs": [],
-   "execution_count": null
-  },
-  {
-   "cell_type": "code",
-   "source": "\n# ── Save LoRA weights + upload everything to HF ──────────────\nfrom huggingface_hub import HfApi, login\nfrom kaggle_secrets import UserSecretsClient\n\nlogin(token=UserSecretsClient().get_secret(\"HF_TOKEN\"))\napi = HfApi()\nREPO = \"CoBeDigger/pipelineops-arena-llama3-grpo\"\n\n# Save LoRA adapter weights locally\nmodel.save_pretrained(\"grpo_lora_adapter\")\ntokenizer.save_pretrained(\"grpo_lora_adapter\")\nprint(\"Saved LoRA adapter locally\")\n\n# Upload adapter weights\napi.upload_folder(folder_path=\"grpo_lora_adapter\", repo_id=REPO, path_in_repo=\"grpo_lora_adapter\")\nprint(\"Uploaded LoRA adapter to HF\")\n\n# Upload training curves\napi.upload_file(path_or_fileobj=\"grpo_training_curve.png\", path_in_repo=\"grpo_training_curve.png\", repo_id=REPO)\napi.upload_file(path_or_fileobj=\"grpo_rewards.json\",       path_in_repo=\"grpo_rewards.json\",       repo_id=REPO)\nprint(\"Uploaded training curves\")\n\n# Summary\nfirst5 = sum(grpo_rewards[:5]) / 5\nlast5  = sum(grpo_rewards[-5:]) / 5\nprint(f\"\\n📊 GRPO Training Summary:\")\nprint(f\"   Steps:              {TRAIN_STEPS}\")\nprint(f\"   Group size (G):     {G}\")\nprint(f\"   First 5 avg reward: {first5:.3f}\")\nprint(f\"   Last  5 avg reward: {last5:.3f}\")\nprint(f\"   Improvement:        +{last5 - first5:.3f}\")\nprint(f\"\\n🔗 https://huggingface.co/{REPO}\")\n",
-   "metadata": {},
-   "execution_count": null,
-   "outputs": []
-  }
- ]
-}