changes made

environment.py

@@ -12,7 +12,7 @@ from dataclasses import dataclass, field
 
 from tasks import TASKS, Task
 from simulator import Simulator
-from rewards import compute_reward, RewardBreakdown
+from rewards import compute_reward
 
 
 @dataclass
@@ -54,27 +54,31 @@ class SWEbenchINEnvironment:
     def _start_container(self) -> str:
         """Start a new Docker container from the swebench-in image."""
         try:
-            result = subprocess.run(
+            run_result = subprocess.run(
                 ["docker", "run", "-d", "--name", "swebench-in-env",
                  "-p", "8080:8080", "swebench-in"],
                 capture_output=True,
                 text=True,
                 timeout=30,
             )
-            container_id = result.stdout.strip()
-            if not container_id:
-                # If container already exists, use it
-                result = subprocess.run(
+            container_id = run_result.stdout.strip()
+            if run_result.returncode == 0 and container_id:
+                return container_id
+
+            # If container already exists (or run failed), try starting it.
+            start_result = subprocess.run(
                     ["docker", "start", "swebench-in-env"],
                     capture_output=True,
                     text=True,
                     timeout=10,
                 )
+            if start_result.returncode == 0:
                 return "swebench-in-env"
-            return container_id
         except (subprocess.TimeoutExpired, FileNotFoundError):
-            # Fallback: return a placeholder for demo/testing without Docker
-            return "swebench-in-env"
+            pass
+
+        # Fallback: return a placeholder for demo/testing without Docker.
+        return "swebench-in-env"
 
     def reset(self, task_id: int = None) -> str:
         """

notebooks/training.ipynb

@@ -1,158 +1,84 @@
 {
  "cells": [
   {
-   "cell_type": "markdown",
-   "metadata": {},
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "_cell_guid": "04149943-a415-4b34-9912-b0f8c528bb41",
+    "_uuid": "8062a470-1b78-44c2-8d72-c8d7b56a6646",
+    "collapsed": false,
+    "jupyter": {
+     "outputs_hidden": false
+    },
+    "trusted": true
+   },
+   "outputs": [],
    "source": [
-    "# SWEbench-IN — GRPO Training Notebook\n",
-    "\n",
-    "This notebook trains a Qwen2.5-3B-Instruct model using GRPO (Group Relative Policy Optimization)\n",
-    "to act as an Indian SWE — fixing broken Linux systems while managing stakeholder communication.\n",
-    "\n",
-    "**Prerequisites:**\n",
-    "- A running SWEbench-IN HuggingFace Space\n",
-    "- A Weights & Biases account\n",
-    "- Google Colab with GPU runtime (T4 or better)"
+    "!pip install -q \\\n",
+    "  \"huggingface_hub>=0.24.0\" \\\n",
+    "  \"unsloth\" \\\n",
+    "  \"unsloth_zoo\"\n",
+    "!pip install \"unsloth[kaggle-new] @ git+https://github.com/unslothai/unsloth.git\" -q\n",
+    "!pip install trl transformers accelerate datasets wandb requests matplotlib pandas -q"
    ]
   },
   {
-   "cell_type": "markdown",
-   "metadata": {},
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "_cell_guid": "a3cb3b13-01ca-4099-89e0-a4bf48071b56",
+    "_uuid": "96b2917f-1eff-408f-a4c6-f24b2b2dca11",
+    "collapsed": false,
+    "jupyter": {
+     "outputs_hidden": false
+    },
+    "trusted": true
+   },
+   "outputs": [],
    "source": [
-    "## Cell 1 — Install Dependencies"
+    "# Run this cell immediately after install finishes\n",
+    "#import os\n",
+    "#os.kill(os.getpid(), 9)"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {},
+   "metadata": {
+    "_cell_guid": "439072a1-3197-4823-9132-995b64208462",
+    "_uuid": "07eb3e2d-de23-4107-867e-651bf08a6915",
+    "collapsed": false,
+    "jupyter": {
+     "outputs_hidden": false
+    },
+    "trusted": true
+   },
    "outputs": [],
    "source": [
+    "import torch\n",
+    "print(\"torch:\", torch.__version__)\n",
+    "print(\"CUDA:\", torch.cuda.is_available())\n",
+    "print(\"GPUs:\", torch.cuda.device_count())\n",
     "\n",
-    "!pip install \"unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git\" -q\n",
-    "!pip install trl transformers accelerate datasets wandb openenv-core -q\n",
-    "!pip install matplotlib pandas -q"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Cell 2 — Import and Configure"
+    "import unsloth\n",
+    "print(\"unsloth:\", unsloth.__version__)\n",
+    "\n",
+    "print(\"All good - ready to train!\")"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m If you're specifying your api key in code, ensure this code is not shared publicly.\n",
-      "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m Consider setting the WANDB_API_KEY environment variable, or running `wandb login` from the command line.\n",
-      "\u001b[34m\u001b[1mwandb\u001b[0m: [wandb.login()] Using explicit session credentials for https://api.wandb.ai.\n",
-      "\u001b[34m\u001b[1mwandb\u001b[0m: No netrc file found, creating one.\n",
-      "\u001b[34m\u001b[1mwandb\u001b[0m: Appending key for api.wandb.ai to your netrc file: /root/.netrc\n",
-      "\u001b[34m\u001b[1mwandb\u001b[0m: Currently logged in as: \u001b[33myusufindian09\u001b[0m (\u001b[33myusufindian09-aaa\u001b[0m) to \u001b[32mhttps://api.wandb.ai\u001b[0m. Use \u001b[1m`wandb login --relogin`\u001b[0m to force relogin\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "🦥 Unsloth: Will patch your computer to enable 2x faster free finetuning.\n",
-      "🦥 Unsloth Zoo will now patch everything to make training faster!\n"
-     ]
-    },
-    {
-     "data": {
-      "text/html": [],
-      "text/plain": [
-       "<IPython.core.display.HTML object>"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "text/html": [
-       "Tracking run with wandb version 0.25.1"
-      ],
-      "text/plain": [
-       "<IPython.core.display.HTML object>"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "text/html": [
-       "Run data is saved locally in <code>/content/wandb/run-20260425_092404-bjgtcv44</code>"
-      ],
-      "text/plain": [
-       "<IPython.core.display.HTML object>"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "text/html": [
-       "Syncing run <strong><a href='https://wandb.ai/yusufindian09-aaa/swebench-in/runs/bjgtcv44' target=\"_blank\">grpo-run-1</a></strong> to <a href='https://wandb.ai/yusufindian09-aaa/swebench-in' target=\"_blank\">Weights & Biases</a> (<a href='https://wandb.me/developer-guide' target=\"_blank\">docs</a>)<br>"
-      ],
-      "text/plain": [
-       "<IPython.core.display.HTML object>"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "text/html": [
-       " View project at <a href='https://wandb.ai/yusufindian09-aaa/swebench-in' target=\"_blank\">https://wandb.ai/yusufindian09-aaa/swebench-in</a>"
-      ],
-      "text/plain": [
-       "<IPython.core.display.HTML object>"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "text/html": [
-       " View run at <a href='https://wandb.ai/yusufindian09-aaa/swebench-in/runs/bjgtcv44' target=\"_blank\">https://wandb.ai/yusufindian09-aaa/swebench-in/runs/bjgtcv44</a>"
-      ],
-      "text/plain": [
-       "<IPython.core.display.HTML object>"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "wandb: Detected [huggingface_hub.inference, mcp, openai] in use.\n",
-      "wandb: Use W&B Weave for improved LLM call tracing. Install Weave with `pip install weave` then add `import weave` to the top of your script.\n",
-      "wandb: For more information, check out the docs at: https://weave-docs.wandb.ai/\n"
-     ]
+   "execution_count": null,
+   "metadata": {
+    "_cell_guid": "a967579b-40a3-4938-aa62-0cafb31ae8d6",
+    "_uuid": "6879acb6-3850-4ae8-ad82-fcac574fa422",
+    "collapsed": false,
+    "jupyter": {
+     "outputs_hidden": false
     },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Setup complete\n"
-     ]
-    }
-   ],
+    "trusted": true
+   },
+   "outputs": [],
    "source": [
     "import os\n",
     "import json\n",
@@ -162,40 +88,69 @@
     "import matplotlib.pyplot as plt\n",
     "import pandas as pd\n",
     "import wandb\n",
-    "wandb.login(key=\"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\")\n",
+    "import requests\n",
+    "\n",
     "from datasets import Dataset\n",
     "from unsloth import FastLanguageModel\n",
     "from trl import GRPOTrainer, GRPOConfig\n",
     "\n",
-    "HF_SPACE_URL      = \"https://huggingface.co/spaces/YUS200619/swebench-ind\"\n",
+    "# ── YOUR CONFIG ──────────────────────────\n",
+    "HF_SPACE_URL      = \"https://YUS200619-swebench-ind.hf.space\"\n",
+    "WANDB_API_KEY     = \"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\"\n",
     "WANDB_PROJECT     = \"swebench-in\"\n",
     "MODEL_NAME        = \"unsloth/Qwen2.5-3B-Instruct\"\n",
     "MAX_SEQ_LEN       = 2048\n",
     "MAX_STEPS         = 15\n",
     "BASELINE_EPISODES = 20\n",
+    "# ─────────────────────────────────────────\n",
     "\n",
+    "os.environ[\"WANDB_API_KEY\"] = WANDB_API_KEY\n",
     "wandb.init(project=WANDB_PROJECT, name=\"grpo-run-1\")\n",
-    "print(\"Setup complete\")"
+    "print(\"Wandb initialized\")"
    ]
   },
   {
-   "cell_type": "markdown",
-   "metadata": {},
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "_cell_guid": "2a5bf235-85f4-4b82-84b8-25793d6d109d",
+    "_uuid": "96cc953e-eba7-4272-bb8d-f712f69fa805",
+    "collapsed": false,
+    "jupyter": {
+     "outputs_hidden": false
+    },
+    "trusted": true
+   },
+   "outputs": [],
    "source": [
-    "## Cell 3 — Load Model (Qwen2.5-3B-Instruct, 4-bit QLoRA via Unsloth)"
+    "try:\n",
+    "    resp = requests.get(f\"{HF_SPACE_URL}/health\", timeout=15)\n",
+    "    print(\"Environment status:\", resp.json())\n",
+    "    print(\"Environment is READY\")\n",
+    "except Exception as e:\n",
+    "    print(\"ENVIRONMENT NOT REACHABLE:\", e)\n",
+    "    print(\"STOP HERE. Fix your HF Space before continuing.\")"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {},
+   "metadata": {
+    "_cell_guid": "e9410e62-1f4c-4ac7-9155-15db5dabcffd",
+    "_uuid": "feb7ae57-5df5-43f0-a057-bb8bd2e9cd8a",
+    "collapsed": false,
+    "jupyter": {
+     "outputs_hidden": false
+    },
+    "trusted": true
+   },
    "outputs": [],
    "source": [
     "model, tokenizer = FastLanguageModel.from_pretrained(\n",
     "    model_name=MODEL_NAME,\n",
     "    max_seq_length=MAX_SEQ_LEN,\n",
     "    dtype=None,\n",
-    "    load_in_4bit=True,       # critical for Colab GPU memory\n",
+    "    load_in_4bit=True,\n",
     ")\n",
     "\n",
     "model = FastLanguageModel.get_peft_model(\n",
@@ -205,40 +160,32 @@
     "    lora_alpha=16,\n",
     "    lora_dropout=0,\n",
     "    bias=\"none\",\n",
-    "    use_gradient_checkpointing=\"unsloth\",  # saves even more memory\n",
+    "    use_gradient_checkpointing=\"unsloth\",\n",
     "    random_state=42,\n",
     ")\n",
     "\n",
-    "print(f\"Model loaded: {MODEL_NAME}\")\n",
-    "print(f\"Trainable parameters: {sum(p.numel() for p in model.parameters() if p.requires_grad):,}\")"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Cell 4 — Define Rollout Function"
+    "print(\"Model loaded successfully\")\n",
+    "print(f\"Trainable params: {sum(p.numel() for p in model.parameters() if p.requires_grad):,}\")"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {},
+   "metadata": {
+    "_cell_guid": "2a55e605-0401-4f0f-87af-9c56bc85c562",
+    "_uuid": "10a8c794-e27a-4351-adf6-b25cb61c0b84",
+    "collapsed": false,
+    "jupyter": {
+     "outputs_hidden": false
+    },
+    "trusted": true
+   },
    "outputs": [],
    "source": [
-    "import requests\n",
-    "\n",
     "class SWEBenchINClient:\n",
-    "    \"\"\"\n",
-    "    Simple HTTP client for your HF Space environment.\n",
-    "    Calls the OpenEnv-compliant API endpoints.\n",
-    "    \"\"\"\n",
     "    def __init__(self, base_url: str):\n",
     "        self.base_url = base_url.rstrip(\"/\")\n",
     "        self.session = requests.Session()\n",
-    "        # Test connection\n",
-    "        resp = self.session.get(f\"{self.base_url}/health\", timeout=10)\n",
-    "        print(f\"Environment connected: {resp.json()}\")\n",
     "\n",
     "    def reset(self, task_id: int = None) -> dict:\n",
     "        payload = {\"task_id\": task_id} if task_id else {}\n",
@@ -266,33 +213,39 @@
     "        )\n",
     "\n",
     "    def state(self) -> dict:\n",
-    "        resp = self.session.get(f\"{self.base_url}/state\", timeout=10)\n",
+    "        resp = self.session.get(\n",
+    "            f\"{self.base_url}/state\",\n",
+    "            timeout=10\n",
+    "        )\n",
     "        return resp.json()\n",
     "\n",
     "    def grade(self) -> dict:\n",
-    "        resp = self.session.post(f\"{self.base_url}/grade\", timeout=30)\n",
+    "        resp = self.session.post(\n",
+    "            f\"{self.base_url}/grade\",\n",
+    "            timeout=30\n",
+    "        )\n",
     "        return resp.json()\n",
     "\n",
+    "env = SWEBenchINClient(HF_SPACE_URL)\n",
     "\n",
-    "env = SWEBenchINClient(https://huggingface.co/spaces/YUS200619/swebench-ind)"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Cell 5 — Curriculum Training Loop\n",
-    "\n",
-    "Curriculum escalates when average reward over last 50 episodes crosses 0.6:\n",
-    "- **Tier 1** (Steps 0–200): Tasks 1+2 only (easy, technical reward)\n",
-    "- **Tier 2** (Steps 200–500): Add Tasks 3+4 (communication reward added)\n",
-    "- **Tier 3** (Steps 500+): Add Task 5 (leave protection added)"
+    "# Quick test\n",
+    "obs = env.reset(task_id=1)\n",
+    "print(\"Reset works:\", type(obs))\n",
+    "print(\"Observation keys:\", list(obs.keys()))"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {},
+   "metadata": {
+    "_cell_guid": "2725c6ef-a52f-4f6f-9c02-42f92cab4b58",
+    "_uuid": "be8a2067-3dc6-41cb-94c5-99d7854b7a5c",
+    "collapsed": false,
+    "jupyter": {
+     "outputs_hidden": false
+    },
+    "trusted": true
+   },
    "outputs": [],
    "source": [
     "SYSTEM_PROMPT = \"\"\"You are an Indian software engineer working on a Linux server.\n",
@@ -301,7 +254,7 @@
     "\n",
     "RULES:\n",
     "- Never use sudo\n",
-    "- Never use rm -rf  \n",
+    "- Never use rm -rf\n",
     "- Fix the server/code issue first\n",
     "- Then reply to manager/client if messages exist\n",
     "- Be efficient — fewer actions is better\n",
@@ -313,25 +266,18 @@
     "{\"type\": \"run_tests\", \"args\": \"\"}\n",
     "{\"type\": \"check_server\", \"args\": \"\"}\n",
     "{\"type\": \"reply_slack\", \"args\": \"Server is back up. Fixed the missing dependency.\"}\n",
-    "{\"type\": \"reply_email\", \"args\": \"Apologies for the downtime. Issue resolved at 2:15 PM.\"}\n",
+    "{\"type\": \"reply_email\", \"args\": \"Apologies for the downtime. Issue resolved.\"}\n",
     "{\"type\": \"close_case\", \"args\": \"\"}\n",
     "\n",
     "Output ONLY valid JSON. Nothing else.\"\"\"\n",
     "\n",
     "\n",
     "def parse_action(text: str) -> dict:\n",
-    "    \"\"\"\n",
-    "    Parse model output into action dict.\n",
-    "    Tries JSON first, then regex, then safe default.\n",
-    "    \"\"\"\n",
     "    text = text.strip()\n",
-    "\n",
-    "    # Strip markdown code blocks if present\n",
     "    text = re.sub(r\"```json\\s*\", \"\", text)\n",
     "    text = re.sub(r\"```\\s*\", \"\", text)\n",
     "    text = text.strip()\n",
     "\n",
-    "    # Try JSON parse\n",
     "    try:\n",
     "        action = json.loads(text)\n",
     "        if \"type\" in action:\n",
@@ -339,21 +285,20 @@
     "    except json.JSONDecodeError:\n",
     "        pass\n",
     "\n",
-    "    # Try regex for known action types\n",
     "    pattern = r'\"?type\"?\\s*[:=]\\s*\"?(\\w+)\"?'\n",
     "    match = re.search(pattern, text)\n",
     "    if match:\n",
     "        action_type = match.group(1)\n",
-    "        args_match = re.search(r'\"?args\"?\\s*[:=]\\s*\"?([^\"}\\n]+)\"?', text)\n",
+    "        args_match = re.search(\n",
+    "            r'\"?args\"?\\s*[:=]\\s*\"?([^\"}\\n]+)\"?', text\n",
+    "        )\n",
     "        args = args_match.group(1).strip() if args_match else \"\"\n",
     "        return {\"type\": action_type, \"args\": args}\n",
     "\n",
-    "    # Safe fallback\n",
     "    return {\"type\": \"close_case\", \"args\": \"\"}\n",
     "\n",
     "\n",
     "def format_prompt(observation: dict) -> str:\n",
-    "    \"\"\"Format the environment observation into a prompt.\"\"\"\n",
     "    return f\"\"\"CURRENT SITUATION:\n",
     "{json.dumps(observation, indent=2)}\n",
     "\n",
@@ -361,10 +306,6 @@
     "\n",
     "\n",
     "def run_episode(task_id: int = None, temperature: float = 0.7) -> tuple:\n",
-    "    \"\"\"\n",
-    "    Run one full episode.\n",
-    "    Returns (list_of_actions, total_reward, final_grade)\n",
-    "    \"\"\"\n",
     "    observation = env.reset(task_id=task_id)\n",
     "    actions_taken = []\n",
     "    total_reward = 0.0\n",
@@ -376,7 +317,6 @@
     "    ]\n",
     "\n",
     "    for step_num in range(MAX_STEPS):\n",
-    "        # Format input for model\n",
     "        input_text = tokenizer.apply_chat_template(\n",
     "            messages,\n",
     "            tokenize=False,\n",
@@ -390,128 +330,82 @@
     "            max_length=MAX_SEQ_LEN\n",
     "        ).to(model.device)\n",
     "\n",
-    "        # Generate action\n",
-    "        with model.disable_adapter() if temperature == 0 else __import__('contextlib').nullcontext():\n",
-    "            outputs = model.generate(\n",
-    "                **inputs,\n",
-    "                max_new_tokens=150,\n",
-    "                temperature=temperature,\n",
-    "                do_sample=(temperature > 0),\n",
-    "                pad_token_id=tokenizer.eos_token_id,\n",
-    "            )\n",
+    "        outputs = model.generate(\n",
+    "            **inputs,\n",
+    "            max_new_tokens=150,\n",
+    "            temperature=max(temperature, 0.01),\n",
+    "            do_sample=True,\n",
+    "            pad_token_id=tokenizer.eos_token_id,\n",
+    "        )\n",
     "\n",
-    "        # Decode only new tokens\n",
     "        new_tokens = outputs[0][inputs[\"input_ids\"].shape[1]:]\n",
-    "        action_text = tokenizer.decode(new_tokens, skip_special_tokens=True)\n",
+    "        action_text = tokenizer.decode(\n",
+    "            new_tokens,\n",
+    "            skip_special_tokens=True\n",
+    "        )\n",
     "\n",
-    "        # Parse action\n",
     "        action = parse_action(action_text)\n",
     "        actions_taken.append({\n",
     "            \"step\": step_num,\n",
-    "            \"raw_output\": action_text[:200],\n",
+    "            \"raw\": action_text[:200],\n",
     "            \"parsed\": action\n",
     "        })\n",
     "\n",
-    "        # Step environment\n",
     "        try:\n",
     "            observation, reward, done, info = env.step(action)\n",
     "            total_reward += reward\n",
     "        except Exception as e:\n",
-    "            print(f\"  Step {step_num} error: {e}\")\n",
+    "            print(f\"  Step error: {e}\")\n",
     "            total_reward -= 0.1\n",
     "            done = True\n",
     "            break\n",
     "\n",
-    "        # Add to conversation history\n",
     "        messages.append({\"role\": \"assistant\", \"content\": action_text})\n",
     "        messages.append({\n",
     "            \"role\": \"user\",\n",
-    "            \"content\": f\"Result: {json.dumps(observation, indent=2)}\\n\\nNext action?\"\n",
+    "            \"content\": f\"Result:\\n{json.dumps(observation, indent=2)}\\n\\nNext action?\"\n",
     "        })\n",
     "\n",
     "        if done:\n",
     "            break\n",
     "\n",
-    "    # Get final grade\n",
     "    try:\n",
     "        final_grade = env.grade()\n",
     "    except:\n",
     "        final_grade = {\"total\": total_reward}\n",
     "\n",
-    "    return actions_taken, total_reward, final_grade"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Cell 6 — Save Model Correctly\n",
+    "    return actions_taken, total_reward, final_grade\n",
     "\n",
-    "**CRITICAL:** Do NOT merge LoRA into 4-bit base model — this damages quality.\n",
-    "Use `save_pretrained` with `method=\"lora\"`."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "print(\"=\" * 50)\n",
-    "print(\"MEASURING UNTRAINED BASELINE\")\n",
-    "print(\"=\" * 50)\n",
-    "\n",
-    "baseline_rewards = []\n",
-    "baseline_task_rewards = {1: [], 2: [], 3: [], 4: [], 5: []}\n",
     "\n",
-    "for i in range(BASELINE_EPISODES):\n",
-    "    task_id = random.choice([1, 2, 3, 4, 5])\n",
-    "    _, reward, grade = run_episode(task_id=task_id, temperature=0.0)\n",
-    "    baseline_rewards.append(reward)\n",
-    "    baseline_task_rewards[task_id].append(reward)\n",
-    "    print(f\"  Episode {i+1:02d}/20 | task={task_id} | reward={reward:.3f}\")\n",
-    "\n",
-    "REAL_BASELINE_AVG = np.mean(baseline_rewards)\n",
-    "print(f\"\\nReal baseline average: {REAL_BASELINE_AVG:.3f}\")\n",
-    "print(f\"Baseline std: {np.std(baseline_rewards):.3f}\")\n",
-    "print(\"Save this number. You need it for the final plot.\")"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Cell 7 — Generate and Commit Training Plots\n",
+    "print(\"System prompt and rollout function ready\")\n",
     "\n",
-    "Both plots must be committed as `.png` files to the repo.\n",
-    "Wandb-only links do not count for the automated validation check."
+    "# Quick sanity check — run one episode\n",
+    "print(\"\\nRunning one test episode...\")\n",
+    "actions, reward, grade = run_episode(task_id=1, temperature=0.0)\n",
+    "print(f\"Test episode reward: {reward:.3f}\")\n",
+    "print(f\"Actions taken: {len(actions)}\")\n",
+    "print(f\"Grade: {grade}\")"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {},
+   "metadata": {
+    "_cell_guid": "2eb94693-a98f-49d9-a81d-75c8cf645b35",
+    "_uuid": "b111970f-7c82-4137-9223-f4da7bd12916",
+    "collapsed": false,
+    "jupyter": {
+     "outputs_hidden": false
+    },
+    "trusted": true
+   },
    "outputs": [],
    "source": [
-    "\"\"\"\n",
-    "GRPO needs a dataset of prompts to generate completions from.\n",
-    "Each prompt = one episode starting state.\n",
-    "We generate 200 prompts across all tasks with curriculum weighting.\n",
-    "\"\"\"\n",
-    "\n",
     "def generate_prompt_dataset(n_prompts: int = 200) -> Dataset:\n",
     "    prompts = []\n",
-    "\n",
-    "    # Curriculum weighting — more easy tasks early\n",
-    "    task_weights = {\n",
-    "        1: 0.30,   # easy\n",
-    "        2: 0.30,   # easy\n",
-    "        3: 0.20,   # medium\n",
-    "        4: 0.15,   # medium\n",
-    "        5: 0.05,   # hard\n",
-    "    }\n",
-    "\n",
+    "    task_weights = {1: 0.30, 2: 0.30, 3: 0.20, 4: 0.15, 5: 0.05}\n",
     "    task_pool = []\n",
+    "\n",
     "    for task_id, weight in task_weights.items():\n",
     "        count = int(n_prompts * weight)\n",
     "        task_pool.extend([task_id] * count)\n",
@@ -519,153 +413,113 @@
     "    random.shuffle(task_pool)\n",
     "\n",
     "    for task_id in task_pool:\n",
-    "        obs = env.reset(task_id=task_id)\n",
-    "        prompt = tokenizer.apply_chat_template(\n",
-    "            [\n",
-    "                {\"role\": \"system\", \"content\": SYSTEM_PROMPT},\n",
-    "                {\"role\": \"user\", \"content\": format_prompt(obs)}\n",
-    "            ],\n",
-    "            tokenize=False,\n",
-    "            add_generation_prompt=True\n",
-    "        )\n",
-    "        prompts.append({\n",
-    "            \"prompt\": prompt,\n",
-    "            \"task_id\": task_id\n",
-    "        })\n",
+    "        try:\n",
+    "            obs = env.reset(task_id=task_id)\n",
+    "            prompt = tokenizer.apply_chat_template(\n",
+    "                [\n",
+    "                    {\"role\": \"system\", \"content\": SYSTEM_PROMPT},\n",
+    "                    {\"role\": \"user\", \"content\": format_prompt(obs)}\n",
+    "                ],\n",
+    "                tokenize=False,\n",
+    "                add_generation_prompt=True\n",
+    "            )\n",
+    "            prompts.append({\"prompt\": prompt, \"task_id\": task_id})\n",
+    "        except Exception as e:\n",
+    "            print(f\"Skipping task {task_id}: {e}\")\n",
+    "            continue\n",
     "\n",
+    "    print(f\"Generated {len(prompts)} prompts\")\n",
     "    return Dataset.from_list(prompts)\n",
     "\n",
     "\n",
-    "print(\"Generating training prompts...\")\n",
     "train_dataset = generate_prompt_dataset(n_prompts=200)\n",
-    "print(f\"Dataset ready: {len(train_dataset)} prompts\")\n"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "code 8"
+    "print(f\"Dataset ready: {len(train_dataset)} prompts\")\n",
+    "print(f\"Sample prompt length: {len(train_dataset[0]['prompt'])} chars\")"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {},
+   "metadata": {
+    "_cell_guid": "7bd8e434-930b-4391-93f4-41a8ee3c9971",
+    "_uuid": "689da08e-d727-4e7b-85fb-009b366e3f50",
+    "collapsed": false,
+    "jupyter": {
+     "outputs_hidden": false
+    },
+    "trusted": true
+   },
    "outputs": [],
    "source": [
-    "\"\"\"\n",
-    "This is what GRPO calls to score each completion.\n",
-    "GRPO generates multiple completions per prompt, scores them,\n",
-    "and updates weights toward higher-scoring completions.\n",
-    "\"\"\"\n",
-    "\n",
-    "episode_log = []  # track all episodes during training\n",
-    "\n",
-    "def grpo_reward_function(completions: list, prompts: list, **kwargs) -> list[float]:\n",
-    "    \"\"\"\n",
-    "    Called by GRPOTrainer after generating completions.\n",
-    "    Returns a reward score for each completion.\n",
-    "    \"\"\"\n",
+    "episode_log = []\n",
+    "\n",
+    "def grpo_reward_function(\n",
+    "    completions: list,\n",
+    "    prompts: list,\n",
+    "    **kwargs\n",
+    ") -> list:\n",
+    "\n",
     "    rewards = []\n",
     "\n",
-    "    for i, completion in enumerate(completions):\n",
+    "    for completion in completions:\n",
     "        try:\n",
-    "            # Parse the action\n",
     "            action = parse_action(completion)\n",
-    "\n",
-    "            # Reset env and take the action\n",
-    "            task_id = random.choice([1, 2, 3, 4])  # avoid task 5 early\n",
-    "            obs = env.reset(task_id=task_id)\n",
+    "            task_id = random.choice([1, 2, 3, 4])\n",
+    "            env.reset(task_id=task_id)\n",
     "            _, reward, done, _ = env.step(action)\n",
     "\n",
-    "            # Continue episode for up to 3 more steps\n",
-    "            # (gives richer signal than single-step reward)\n",
-    "            for _ in range(3):\n",
-    "                if done:\n",
-    "                    break\n",
-    "                # Simple follow-up: try to close case\n",
-    "                followup = {\"type\": \"check_server\", \"args\": \"\"}\n",
-    "                obs, step_reward, done, _ = env.step(followup)\n",
-    "                reward += step_reward * 0.5  # discount follow-up rewards\n",
-    "\n",
-    "            # Penalty for bad actions\n",
     "            if action[\"type\"] == \"close_case\" and reward < 0.1:\n",
-    "                reward -= 0.3  # penalize giving up immediately\n",
+    "                reward -= 0.3\n",
     "\n",
     "            rewards.append(float(reward))\n",
     "\n",
-    "        except Exception as e:\n",
-    "            # Environment error — penalize\n",
+    "        except Exception:\n",
     "            rewards.append(-0.5)\n",
     "\n",
-    "    # Log to wandb\n",
     "    wandb.log({\n",
     "        \"reward/batch_mean\": np.mean(rewards),\n",
-    "        \"reward/batch_std\": np.std(rewards),\n",
-    "        \"reward/batch_max\": np.max(rewards),\n",
+    "        \"reward/batch_max\":  np.max(rewards),\n",
+    "        \"reward/batch_std\":  np.std(rewards),\n",
     "    })\n",
     "\n",
     "    episode_log.extend(rewards)\n",
-    "    return rewards\n"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "code 9\n"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "training_args = GRPOConfig(\n",
-    "    # Output\n",
-    "    output_dir=\"./swebench-in-checkpoints\",\n",
+    "    return rewards\n",
     "\n",
-    "    # Training duration\n",
-    "    num_train_epochs=3,\n",
-    "    max_steps=500,              # hard cap — enough for hackathon\n",
     "\n",
-    "    # Batch sizes — small to fit T4 GPU\n",
+    "training_args = GRPOConfig(\n",
+    "    output_dir=\"./swebench-checkpoints\",\n",
+    "    num_train_epochs=3,\n",
+    "    max_steps=500,\n",
     "    per_device_train_batch_size=2,\n",
-    "    gradient_accumulation_steps=8,  # effective batch = 16\n",
-    "\n",
-    "    # GRPO specific\n",
-    "    num_generations=4,          # completions per prompt GRPO compares\n",
-    "    max_completion_length=200,  # max tokens per action\n",
-    "\n",
-    "    # Optimizer\n",
-    "    learning_rate=5e-6,         # low LR for RL stability\n",
+    "    gradient_accumulation_steps=8,\n",
+    "    num_generations=4,\n",
+    "    max_completion_length=200,\n",
+    "    learning_rate=5e-6,\n",
     "    warmup_steps=20,\n",
     "    weight_decay=0.01,\n",
-    "\n",
-    "    # Logging\n",
     "    logging_steps=10,\n",
     "    save_steps=100,\n",
     "    report_to=\"wandb\",\n",
-    "\n",
-    "    # Memory\n",
     "    gradient_checkpointing=True,\n",
-    "    fp16=True,                  # use fp16 on T4\n",
-    ")"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "code 10\n"
+    "    bf16=False,\n",
+    "    fp16=True,\n",
+    ")\n",
+    "\n",
+    "print(\"Reward function and training config ready\")"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {},
+   "metadata": {
+    "_cell_guid": "dd675d60-2347-40ce-b612-aac21ac7314f",
+    "_uuid": "21373bdf-191a-4493-b04c-2eade2dd5954",
+    "collapsed": false,
+    "jupyter": {
+     "outputs_hidden": false
+    },
+    "trusted": true
+   },
    "outputs": [],
    "source": [
     "trainer = GRPOTrainer(\n",
@@ -677,27 +531,28 @@
     ")\n",
     "\n",
     "print(\"Starting GRPO training...\")\n",
-    "print(f\"Dataset: {len(train_dataset)} prompts\")\n",
-    "print(f\"Max steps: {training_args.max_steps}\")\n",
+    "print(f\"Steps: {training_args.max_steps}\")\n",
     "print(f\"Generations per prompt: {training_args.num_generations}\")\n",
+    "print(\"This takes 60-90 minutes. Watch wandb dashboard.\")\n",
     "print(\"=\" * 50)\n",
     "\n",
     "trainer.train()\n",
     "\n",
-    "print(\"Training complete!\")\n"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "code 11"
+    "print(\"Training complete!\")"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {},
+   "metadata": {
+    "_cell_guid": "7b5a47d9-7abc-4824-982b-adb2ab684961",
+    "_uuid": "9baa6267-f3ea-49d5-a6b8-de129657fcf9",
+    "collapsed": false,
+    "jupyter": {
+     "outputs_hidden": false
+    },
+    "trusted": true
+   },
    "outputs": [],
    "source": [
     "print(\"=\" * 50)\n",
@@ -708,212 +563,166 @@
     "\n",
     "for i in range(BASELINE_EPISODES):\n",
     "    task_id = random.choice([1, 2, 3, 4, 5])\n",
-    "    _, reward, grade = run_episode(task_id=task_id, temperature=0.0)\n",
+    "    _, reward, _ = run_episode(task_id=task_id, temperature=0.0)\n",
     "    trained_rewards.append(reward)\n",
-    "    print(f\"  Episode {i+1:02d}/20 | task={task_id} | reward={reward:.3f}\")\n",
+    "    print(f\"Episode {i+1:02d}/{BASELINE_EPISODES} \"\n",
+    "          f\"| task={task_id} | reward={reward:.3f}\")\n",
     "\n",
     "TRAINED_AVG = np.mean(trained_rewards)\n",
-    "print(f\"\\nTrained average: {TRAINED_AVG:.3f}\")\n",
+    "print(f\"\\nTrained average:  {TRAINED_AVG:.3f}\")\n",
     "print(f\"Baseline average: {REAL_BASELINE_AVG:.3f}\")\n",
-    "print(f\"Improvement: {TRAINED_AVG - REAL_BASELINE_AVG:.3f}\")"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "code 12"
+    "print(f\"Improvement:      +{TRAINED_AVG - REAL_BASELINE_AVG:.3f}\")"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {},
+   "metadata": {
+    "_cell_guid": "089714bd-53e6-40e7-a041-3b18dbff9f16",
+    "_uuid": "452ccb4d-6b90-4049-a62d-841a054e9bcd",
+    "collapsed": false,
+    "jupyter": {
+     "outputs_hidden": false
+    },
+    "trusted": true
+   },
    "outputs": [],
    "source": [
-    "import os\n",
     "os.makedirs(\"plots\", exist_ok=True)\n",
     "\n",
-    "# Pull training history from wandb\n",
+    "# ── Plot 1: Reward Curve ──────────────────\n",
     "history_df = wandb.run.history(\n",
-    "    keys=[\"reward/batch_mean\", \"_step\", \"loss\"]\n",
+    "    keys=[\"reward/batch_mean\", \"_step\"]\n",
     ")\n",
     "history_df = history_df.dropna(subset=[\"reward/batch_mean\"])\n",
     "\n",
-    "# ── Plot 1: Reward Curve ──────────────────────────────────\n",
     "fig, ax = plt.subplots(figsize=(12, 5))\n",
-    "\n",
-    "# Training reward\n",
     "ax.plot(\n",
     "    history_df[\"_step\"],\n",
     "    history_df[\"reward/batch_mean\"],\n",
-    "    color=\"steelblue\",\n",
-    "    alpha=0.6,\n",
-    "    linewidth=1,\n",
-    "    label=\"Training reward (per batch)\"\n",
+    "    color=\"steelblue\", alpha=0.4,\n",
+    "    linewidth=1, label=\"Per batch reward\"\n",
     ")\n",
-    "\n",
-    "# Smoothed training reward\n",
     "if len(history_df) > 10:\n",
     "    smoothed = pd.Series(\n",
     "        history_df[\"reward/batch_mean\"].values\n",
     "    ).rolling(window=20, min_periods=1).mean()\n",
     "    ax.plot(\n",
-    "        history_df[\"_step\"],\n",
-    "        smoothed,\n",
-    "        color=\"steelblue\",\n",
-    "        linewidth=2.5,\n",
-    "        label=\"Training reward (smoothed)\"\n",
+    "        history_df[\"_step\"], smoothed,\n",
+    "        color=\"steelblue\", linewidth=2.5,\n",
+    "        label=\"Smoothed reward\"\n",
     "    )\n",
-    "\n",
-    "# Baseline and trained horizontal lines\n",
     "ax.axhline(\n",
-    "    y=REAL_BASELINE_AVG,\n",
-    "    color=\"red\",\n",
-    "    linestyle=\"--\",\n",
-    "    linewidth=2,\n",
+    "    y=REAL_BASELINE_AVG, color=\"red\",\n",
+    "    linestyle=\"--\", linewidth=2,\n",
     "    label=f\"Untrained baseline ({REAL_BASELINE_AVG:.2f})\"\n",
     ")\n",
     "ax.axhline(\n",
-    "    y=TRAINED_AVG,\n",
-    "    color=\"green\",\n",
-    "    linestyle=\"--\",\n",
-    "    linewidth=2,\n",
+    "    y=TRAINED_AVG, color=\"green\",\n",
+    "    linestyle=\"--\", linewidth=2,\n",
     "    label=f\"Trained model ({TRAINED_AVG:.2f})\"\n",
     ")\n",
-    "\n",
     "ax.set_xlabel(\"Training Step\", fontsize=12)\n",
     "ax.set_ylabel(\"Episode Reward\", fontsize=12)\n",
-    "ax.set_title(\"SWEbench-IN: GRPO Training Reward Curve\", fontsize=14)\n",
+    "ax.set_title(\"SWEbench-IN: GRPO Training Reward Curve\",\n",
+    "             fontsize=14)\n",
     "ax.legend(fontsize=10)\n",
     "ax.grid(True, alpha=0.3)\n",
     "plt.tight_layout()\n",
-    "plt.savefig(\"plots/reward_curve.png\", dpi=150, bbox_inches=\"tight\")\n",
+    "plt.savefig(\"plots/reward_curve.png\", dpi=150,\n",
+    "            bbox_inches=\"tight\")\n",
     "plt.show()\n",
     "print(\"Saved: plots/reward_curve.png\")\n",
     "\n",
-    "# ── Plot 2: Loss Curve ────────────────────────────────────\n",
+    "# ── Plot 2: Loss Curve ────────────────────\n",
     "fig, ax = plt.subplots(figsize=(12, 5))\n",
-    "\n",
-    "# Try trainer log history first (more reliable than wandb)\n",
     "log_history = trainer.state.log_history\n",
     "losses = [x[\"loss\"] for x in log_history if \"loss\" in x]\n",
     "steps  = [x[\"step\"] for x in log_history if \"loss\" in x]\n",
     "\n",
     "if losses:\n",
-    "    ax.plot(steps, losses, color=\"crimson\", linewidth=1.5, label=\"Policy Loss\")\n",
-    "elif \"loss\" in history_df.columns:\n",
-    "    loss_df = history_df.dropna(subset=[\"loss\"])\n",
-    "    ax.plot(loss_df[\"_step\"], loss_df[\"loss\"],\n",
-    "            color=\"crimson\", linewidth=1.5, label=\"Policy Loss\")\n",
-    "else:\n",
-    "    ax.text(0.5, 0.5, \"Loss not logged\", transform=ax.transAxes,\n",
-    "            ha=\"center\", fontsize=14)\n",
+    "    ax.plot(steps, losses, color=\"crimson\",\n",
+    "            linewidth=1.5, label=\"Policy Loss\")\n",
+    "    ax.set_xlabel(\"Training Step\", fontsize=12)\n",
+    "    ax.set_ylabel(\"Loss\", fontsize=12)\n",
+    "    ax.set_title(\"SWEbench-IN: Policy Loss Curve\",\n",
+    "                 fontsize=14)\n",
+    "    ax.legend(fontsize=10)\n",
+    "    ax.grid(True, alpha=0.3)\n",
     "\n",
-    "ax.set_xlabel(\"Training Step\", fontsize=12)\n",
-    "ax.set_ylabel(\"Loss\", fontsize=12)\n",
-    "ax.set_title(\"SWEbench-IN: Policy Loss Curve\", fontsize=14)\n",
-    "ax.legend(fontsize=10)\n",
-    "ax.grid(True, alpha=0.3)\n",
     "plt.tight_layout()\n",
-    "plt.savefig(\"plots/loss_curve.png\", dpi=150, bbox_inches=\"tight\")\n",
+    "plt.savefig(\"plots/loss_curve.png\", dpi=150,\n",
+    "            bbox_inches=\"tight\")\n",
     "plt.show()\n",
     "print(\"Saved: plots/loss_curve.png\")\n",
     "\n",
-    "# ── Plot 3: Before vs After Comparison ───────────────────\n",
-    "fig, ax = plt.subplots(figsize=(10, 5))\n",
-    "\n",
+    "# ── Plot 3: Before vs After ───────────────\n",
+    "fig, ax = plt.subplots(figsize=(12, 5))\n",
     "episodes = list(range(1, BASELINE_EPISODES + 1))\n",
     "ax.plot(episodes, baseline_rewards,\n",
-    "        color=\"red\", marker=\"o\", linewidth=1.5,\n",
+    "        color=\"red\", marker=\"o\",\n",
+    "        linewidth=1.5,\n",
     "        label=f\"Untrained (avg={REAL_BASELINE_AVG:.2f})\")\n",
     "ax.plot(episodes, trained_rewards,\n",
-    "        color=\"green\", marker=\"s\", linewidth=1.5,\n",
+    "        color=\"green\", marker=\"s\",\n",
+    "        linewidth=1.5,\n",
     "        label=f\"Trained (avg={TRAINED_AVG:.2f})\")\n",
-    "ax.axhline(y=REAL_BASELINE_AVG, color=\"red\",\n",
-    "           linestyle=\"--\", alpha=0.4)\n",
-    "ax.axhline(y=TRAINED_AVG, color=\"green\",\n",
-    "           linestyle=\"--\", alpha=0.4)\n",
-    "ax.fill_between(episodes, baseline_rewards, trained_rewards,\n",
-    "                alpha=0.1, color=\"green\",\n",
-    "                label=f\"Improvement: +{TRAINED_AVG - REAL_BASELINE_AVG:.2f}\")\n",
-    "\n",
+    "ax.fill_between(\n",
+    "    episodes, baseline_rewards, trained_rewards,\n",
+    "    alpha=0.1, color=\"green\",\n",
+    "    label=f\"Improvement: +{TRAINED_AVG - REAL_BASELINE_AVG:.2f}\"\n",
+    ")\n",
     "ax.set_xlabel(\"Episode\", fontsize=12)\n",
     "ax.set_ylabel(\"Reward\", fontsize=12)\n",
-    "ax.set_title(\"SWEbench-IN: Before vs After GRPO Training\", fontsize=14)\n",
+    "ax.set_title(\"SWEbench-IN: Before vs After GRPO Training\",\n",
+    "             fontsize=14)\n",
     "ax.legend(fontsize=10)\n",
     "ax.grid(True, alpha=0.3)\n",
     "plt.tight_layout()\n",
-    "plt.savefig(\"plots/before_after.png\", dpi=150, bbox_inches=\"tight\")\n",
+    "plt.savefig(\"plots/before_after.png\", dpi=150,\n",
+    "            bbox_inches=\"tight\")\n",
     "plt.show()\n",
-    "print(\"Saved: plots/before_after.png\")\n"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "cell 13"
+    "print(\"Saved: plots/before_after.png\")\n",
+    "\n",
+    "print(\"\\nAll plots saved.\")\n",
+    "print(\"Download from Kaggle Files panel → commit to repo\")"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {},
+   "metadata": {
+    "_cell_guid": "16534aa6-897c-453e-9a01-23da3fc20acc",
+    "_uuid": "24c0f45c-3e8b-4205-8ae4-3a503f3624e1",
+    "collapsed": false,
+    "jupyter": {
+     "outputs_hidden": false
+    },
+    "trusted": true
+   },
    "outputs": [],
    "source": [
-    "\"\"\"\n",
-    "CRITICAL: Do NOT merge LoRA into 4-bit base.\n",
-    "Save adapters only using method=\"lora\"\n",
-    "\"\"\"\n",
     "model.save_pretrained_merged(\n",
     "    \"swebench-in-lora\",\n",
     "    tokenizer=tokenizer,\n",
     "    save_method=\"lora\"\n",
     ")\n",
-    "print(\"Model saved to swebench-in-lora/\")\n",
-    "\n",
-    "\n"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "cell 14"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# ============================================================\n",
-    "\"\"\"\n",
-    "AFTER downloading plots from Colab Files panel:\n",
-    "\n",
-    "cd your-local-repo\n",
-    "cp ~/Downloads/reward_curve.png plots/\n",
-    "cp ~/Downloads/loss_curve.png plots/\n",
-    "cp ~/Downloads/before_after.png plots/\n",
-    "\n",
-    "git add plots/\n",
-    "git commit -m \"Add training evidence: reward curve, loss curve, before/after\"\n",
-    "git push origin main\n",
-    "\n",
-    "Then verify from logged-out browser that plots appear in README.\n",
-    "\"\"\"\n",
-    "print(\"Download plots from Colab Files panel → commit to repo\")\n",
-    "print(\"plots/reward_curve.png\")\n",
-    "print(\"plots/loss_curve.png\")\n",
-    "print(\"plots/before_after.png\")"
+    "print(\"Model saved\")"
    ]
   }
  ],
  "metadata": {
+  "kaggle": {
+   "accelerator": "nvidiaTeslaT4",
+   "dataSources": [],
+   "dockerImageVersionId": 31329,
+   "isGpuEnabled": true,
+   "isInternetEnabled": true,
+   "language": "python",
+   "sourceType": "notebook"
+  },
   "kernelspec": {
-   "display_name": "Python 3 (ipykernel)",
+   "display_name": "Python 3",
    "language": "python",
    "name": "python3"
   },
@@ -927,7 +736,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.12.13"
+   "version": "3.12.12"
   }
  },
  "nbformat": 4,