changed notebook

Dockerfile

@@ -20,7 +20,7 @@ COPY openenv.yaml /home/user2/
 COPY requirements.txt /home/user2/
 COPY README.md /home/user2/
 
-EXPOSE 7868
+EXPOSE 7860
 ENV ENABLE_WEB_INTERFACE=true
 CMD ["python", "server/app.py"]
 

notebooks/training.ipynb

@@ -24,11 +24,14 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
-    "%pip install unsloth trl transformers accelerate openenv-core[core]>=0.2.2 wandb -q"
+    "\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"
    ]
   },
   {
@@ -40,9 +43,21 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "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",
@@ -52,37 +67,115 @@
      ]
     },
     {
-     "ename": "ModuleNotFoundError",
-     "evalue": "No module named 'openenv.client'",
-     "output_type": "error",
-     "traceback": [
-      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
-      "\u001b[0;31mModuleNotFoundError\u001b[0m                       Traceback (most recent call last)",
-      "\u001b[0;32m/tmp/ipykernel_1168/3569467806.py\u001b[0m in \u001b[0;36m<cell line: 0>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      6\u001b[0m \u001b[0;32mfrom\u001b[0m \u001b[0munsloth\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0mFastLanguageModel\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      7\u001b[0m \u001b[0;32mfrom\u001b[0m \u001b[0mtrl\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0mGRPOTrainer\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mGRPOConfig\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 8\u001b[0;31m \u001b[0;32mfrom\u001b[0m \u001b[0mopenenv\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mclient\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0mEnvironment\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0mOpenEnvClient\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      9\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     10\u001b[0m \u001b[0mwandb\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0minit\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mproject\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m\"swebench-in\"\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mname\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m\"grpo-run-1\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
-      "\u001b[0;31mModuleNotFoundError\u001b[0m: No module named 'openenv.client'",
-      "",
-      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0;32m\nNOTE: If your import is failing due to a missing package, you can\nmanually install dependencies using either !pip or !apt.\n\nTo view examples of installing some common dependencies, click the\n\"Open Examples\" button below.\n\u001b[0;31m---------------------------------------------------------------------------\u001b[0m\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"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Setup complete\n"
      ]
     }
    ],
    "source": [
     "import os\n",
-    "import wandb\n",
-    "import random\n",
-    "import re\n",
     "import json\n",
+    "import re\n",
+    "import random\n",
+    "import numpy as np\n",
+    "import matplotlib.pyplot as plt\n",
+    "import pandas as pd\n",
+    "import wandb\n",
+    "wandb.login(key=\"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\")\n",
+    "from datasets import Dataset\n",
     "from unsloth import FastLanguageModel\n",
     "from trl import GRPOTrainer, GRPOConfig\n",
-    "from openenv.client import Environment as OpenEnvClient\n",
-    "\n",
-    "wandb.init(project=\"swebench-in\", name=\"grpo-run-1\")\n",
     "\n",
-    "# Set this in Colab via: %env HF_SPACE_URL=https://huggingface.co/spaces/YUS200619/swebench-ind\n",
-    "HF_SPACE_URL = os.getenv(\"HF_SPACE_URL\", \"\").strip()\n",
-    "if not HF_SPACE_URL:\n",
-    "    raise ValueError(\"HF_SPACE_URL is empty. Set it before running this cell.\")\n",
+    "HF_SPACE_URL      = \"https://huggingface.co/spaces/YUS200619/swebench-ind\"\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",
-    "env = OpenEnvClient(HF_SPACE_URL)"
+    "wandb.init(project=WANDB_PROJECT, name=\"grpo-run-1\")\n",
+    "print(\"Setup complete\")"
    ]
   },
   {
@@ -99,11 +192,12 @@
    "outputs": [],
    "source": [
     "model, tokenizer = FastLanguageModel.from_pretrained(\n",
-    "    model_name=\"Qwen/Qwen2.5-3B-Instruct\",\n",
-    "    max_seq_length=2048,\n",
+    "    model_name=MODEL_NAME,\n",
+    "    max_seq_length=MAX_SEQ_LEN,\n",
     "    dtype=None,\n",
-    "    load_in_4bit=True,\n",
+    "    load_in_4bit=True,       # critical for Colab GPU memory\n",
     ")\n",
+    "\n",
     "model = FastLanguageModel.get_peft_model(\n",
     "    model,\n",
     "    r=16,\n",
@@ -111,8 +205,12 @@
     "    lora_alpha=16,\n",
     "    lora_dropout=0,\n",
     "    bias=\"none\",\n",
-    "    use_gradient_checkpointing=True,\n",
-    ")"
+    "    use_gradient_checkpointing=\"unsloth\",  # saves even more memory\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):,}\")"
    ]
   },
   {
@@ -128,40 +226,55 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "def parse_action(action_text: str) -> dict:\n",
+    "import requests\n",
+    "\n",
+    "class SWEBenchINClient:\n",
     "    \"\"\"\n",
-    "    Parse the model's generated text into an action dict.\n",
-    "    Expected format: ACTION_TYPE: args\n",
+    "    Simple HTTP client for your HF Space environment.\n",
+    "    Calls the OpenEnv-compliant API endpoints.\n",
     "    \"\"\"\n",
-    "    action_text = action_text.strip()\n",
-    "    # Try to find action pattern\n",
-    "    match = re.search(r'(run_command|read_file|write_file|run_tests|check_server|reply_slack|reply_email|reply_hr|close_case)[:\\s]+(.*)', action_text, re.DOTALL)\n",
-    "    if match:\n",
-    "        return {\"type\": match.group(1), \"args\": match.group(2).strip()}\n",
-    "    # Default: treat as run_command\n",
-    "    return {\"type\": \"run_command\", \"args\": action_text}\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",
+    "        resp = self.session.post(\n",
+    "            f\"{self.base_url}/reset\",\n",
+    "            json=payload,\n",
+    "            timeout=30\n",
+    "        )\n",
+    "        resp.raise_for_status()\n",
+    "        return resp.json()\n",
     "\n",
-    "def rollout(prompt: str, task_id: int) -> tuple[list[str], float]:\n",
-    "    \"\"\"\n",
-    "    Run one episode. Return (action_sequence, total_reward).\n",
-    "    Uses sampling with temperature 0.7.\n",
-    "    \"\"\"\n",
-    "    obs = env.reset(task_id=task_id)\n",
-    "    actions = []\n",
-    "    total_reward = 0.0\n",
-    "    done = False\n",
+    "    def step(self, action: dict) -> tuple:\n",
+    "        resp = self.session.post(\n",
+    "            f\"{self.base_url}/step\",\n",
+    "            json={\"action\": action},\n",
+    "            timeout=30\n",
+    "        )\n",
+    "        resp.raise_for_status()\n",
+    "        data = resp.json()\n",
+    "        return (\n",
+    "            data[\"observation\"],\n",
+    "            float(data[\"reward\"]),\n",
+    "            bool(data[\"done\"]),\n",
+    "            data.get(\"info\", {})\n",
+    "        )\n",
     "\n",
-    "    while not done:\n",
-    "        inputs = tokenizer(f\"Observation: {obs}\\nAction:\", return_tensors=\"pt\")\n",
-    "        output = model.generate(**inputs, max_new_tokens=100, do_sample=True, temperature=0.7)\n",
-    "        action_text = tokenizer.decode(output[0], skip_special_tokens=True)\n",
-    "        action = parse_action(action_text)\n",
-    "        obs, reward, done, info = env.step(action)\n",
-    "        actions.append(action_text)\n",
-    "        total_reward += reward\n",
+    "    def state(self) -> dict:\n",
+    "        resp = self.session.get(f\"{self.base_url}/state\", timeout=10)\n",
+    "        return resp.json()\n",
+    "\n",
+    "    def grade(self) -> dict:\n",
+    "        resp = self.session.post(f\"{self.base_url}/grade\", timeout=30)\n",
+    "        return resp.json()\n",
     "\n",
-    "    return actions, total_reward"
+    "\n",
+    "env = SWEBenchINClient(https://huggingface.co/spaces/YUS200619/swebench-ind)"
    ]
   },
   {
@@ -182,45 +295,150 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Curriculum: tier 1 tasks first (1,2), then tier 2 (3,4), then tier 3 (5)\n",
-    "CURRICULUM = {\n",
-    "    \"tier1\": [1, 2],\n",
-    "    \"tier2\": [3, 4],\n",
-    "    \"tier3\": [5],\n",
-    "}\n",
+    "SYSTEM_PROMPT = \"\"\"You are an Indian software engineer working on a Linux server.\n",
+    "You will receive a broken environment description.\n",
+    "Your job: fix the technical issue AND handle any human communication.\n",
     "\n",
-    "current_tier = \"tier1\"\n",
-    "tier_rewards = []\n",
+    "RULES:\n",
+    "- Never use sudo\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",
     "\n",
-    "for step in range(700):\n",
-    "    task_id = random.choice(CURRICULUM[current_tier])\n",
-    "    actions, reward = rollout(\"\", task_id)\n",
+    "OUTPUT FORMAT (pick one per turn):\n",
+    "{\"type\": \"run_command\", \"args\": \"pip install flask\"}\n",
+    "{\"type\": \"read_file\", \"args\": \"/home/user2/app.py\"}\n",
+    "{\"type\": \"write_file\", \"args\": \"/home/user2/app.py\", \"content\": \"fixed code here\"}\n",
+    "{\"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\": \"close_case\", \"args\": \"\"}\n",
     "\n",
-    "    # Log to wandb\n",
-    "    wandb.log({\n",
-    "        \"reward/total\": reward,\n",
-    "        \"training_step\": step,\n",
-    "        \"task_id\": task_id,\n",
-    "        \"current_tier\": current_tier,\n",
-    "        \"num_actions\": len(actions),\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",
+    "            return action\n",
+    "    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 = 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",
-    "    tier_rewards.append(reward)\n",
-    "\n",
-    "    # Escalate curriculum\n",
-    "    if len(tier_rewards) >= 50 and sum(tier_rewards[-50:]) / 50 >= 0.6:\n",
-    "        if current_tier == \"tier1\":\n",
-    "            current_tier = \"tier2\"\n",
-    "            tier_rewards = []\n",
-    "            print(f\"Step {step}: Escalating to tier 2\")\n",
-    "        elif current_tier == \"tier2\":\n",
-    "            current_tier = \"tier3\"\n",
-    "            tier_rewards = []\n",
-    "            print(f\"Step {step}: Escalating to tier 3\")\n",
-    "\n",
-    "    if step % 50 == 0:\n",
-    "        avg = sum(tier_rewards[-50:]) / max(len(tier_rewards[-50:]), 1)\n",
-    "        print(f\"Step {step} | Tier: {current_tier} | Avg reward (last 50): {avg:.3f}\")"
+    "\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",
+    "What is your next action? Output valid JSON only.\"\"\"\n",
+    "\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",
+    "    done = False\n",
+    "\n",
+    "    messages = [\n",
+    "        {\"role\": \"system\", \"content\": SYSTEM_PROMPT},\n",
+    "        {\"role\": \"user\", \"content\": format_prompt(observation)}\n",
+    "    ]\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",
+    "            add_generation_prompt=True\n",
+    "        )\n",
+    "\n",
+    "        inputs = tokenizer(\n",
+    "            input_text,\n",
+    "            return_tensors=\"pt\",\n",
+    "            truncation=True,\n",
+    "            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",
+    "\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",
+    "\n",
+    "        # Parse action\n",
+    "        action = parse_action(action_text)\n",
+    "        actions_taken.append({\n",
+    "            \"step\": step_num,\n",
+    "            \"raw_output\": 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",
+    "            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",
+    "        })\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"
    ]
   },
   {
@@ -239,15 +457,24 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Save LoRA adapters without merging into 4-bit base weights.\n",
-    "if hasattr(model, \"save_pretrained_merged\"):\n",
-    "    model.save_pretrained_merged(\"swebench-in-lora\", tokenizer=tokenizer, save_method=\"lora\")\n",
-    "else:\n",
-    "    model.save_pretrained(\"swebench-in-lora\")\n",
-    "    tokenizer.save_pretrained(\"swebench-in-lora\")\n",
+    "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",
-    "# Optional: push to hub\n",
-    "# model.push_to_hub(\"YOUR_HF_USERNAME/swebench-in-lora\")"
+    "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.\")"
    ]
   },
   {
@@ -266,45 +493,421 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import matplotlib.pyplot as plt\n",
-    "import os\n",
+    "\"\"\"\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_pool = []\n",
+    "    for task_id, weight in task_weights.items():\n",
+    "        count = int(n_prompts * weight)\n",
+    "        task_pool.extend([task_id] * count)\n",
+    "\n",
+    "    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",
+    "\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"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "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",
+    "    rewards = []\n",
+    "\n",
+    "    for i, completion in enumerate(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",
+    "            _, 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",
+    "\n",
+    "            rewards.append(float(reward))\n",
+    "\n",
+    "        except Exception as e:\n",
+    "            # Environment error — penalize\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",
+    "    })\n",
     "\n",
-    "# Pull run history from wandb\n",
-    "run = wandb.run\n",
-    "history = run.history()\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",
+    "\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",
+    "    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",
+    "    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"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "trainer = GRPOTrainer(\n",
+    "    model=model,\n",
+    "    reward_funcs=grpo_reward_function,\n",
+    "    args=training_args,\n",
+    "    train_dataset=train_dataset,\n",
+    "    tokenizer=tokenizer,\n",
+    ")\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\"Generations per prompt: {training_args.num_generations}\")\n",
+    "print(\"=\" * 50)\n",
     "\n",
+    "trainer.train()\n",
+    "\n",
+    "print(\"Training complete!\")\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "code 11"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "print(\"=\" * 50)\n",
+    "print(\"EVALUATING TRAINED MODEL\")\n",
+    "print(\"=\" * 50)\n",
+    "\n",
+    "trained_rewards = []\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",
+    "    trained_rewards.append(reward)\n",
+    "    print(f\"  Episode {i+1:02d}/20 | 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\"Baseline average: {REAL_BASELINE_AVG:.3f}\")\n",
+    "print(f\"Improvement: {TRAINED_AVG - REAL_BASELINE_AVG:.3f}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "code 12"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
     "os.makedirs(\"plots\", exist_ok=True)\n",
     "\n",
-    "# --- Reward Curve ---\n",
-    "fig, ax = plt.subplots(figsize=(10, 5))\n",
-    "ax.plot(history[\"training_step\"], history[\"reward/total\"],\n",
-    "        label=\"Trained Agent\", color=\"steelblue\")\n",
-    "ax.axhline(y=-0.4, color=\"orange\", linestyle=\"--\",\n",
-    "           label=\"Untrained Baseline (-0.4)\")\n",
-    "ax.set_xlabel(\"Training Step\")\n",
-    "ax.set_ylabel(\"Episode Reward\")\n",
-    "ax.set_title(\"SWEbench-IN: Training Reward Curve\")\n",
-    "ax.legend()\n",
+    "# Pull training history from wandb\n",
+    "history_df = wandb.run.history(\n",
+    "    keys=[\"reward/batch_mean\", \"_step\", \"loss\"]\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",
+    ")\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",
+    "    )\n",
+    "\n",
+    "# Baseline and trained horizontal lines\n",
+    "ax.axhline(\n",
+    "    y=REAL_BASELINE_AVG,\n",
+    "    color=\"red\",\n",
+    "    linestyle=\"--\",\n",
+    "    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",
+    "    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.legend(fontsize=10)\n",
     "ax.grid(True, alpha=0.3)\n",
     "plt.tight_layout()\n",
-    "plt.savefig(\"plots/reward_curve.png\", dpi=150)\n",
+    "plt.savefig(\"plots/reward_curve.png\", dpi=150, bbox_inches=\"tight\")\n",
     "plt.show()\n",
-    "print(\"plots/reward_curve.png saved. Commit it to your repo now.\")\n",
+    "print(\"Saved: plots/reward_curve.png\")\n",
+    "\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",
-    "# --- Loss Curve ---\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",
+    "\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.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",
-    "if \"loss\" in history.columns:\n",
-    "    ax.plot(history[\"training_step\"], history[\"loss\"],\n",
-    "            label=\"Policy Loss\", color=\"crimson\")\n",
-    "ax.set_xlabel(\"Training Step\")\n",
-    "ax.set_ylabel(\"Loss\")\n",
-    "ax.set_title(\"SWEbench-IN: Policy Loss Curve\")\n",
-    "ax.legend()\n",
+    "\n",
+    "episodes = list(range(1, BASELINE_EPISODES + 1))\n",
+    "ax.plot(episodes, baseline_rewards,\n",
+    "        color=\"red\", marker=\"o\", 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",
+    "        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.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.legend(fontsize=10)\n",
     "ax.grid(True, alpha=0.3)\n",
     "plt.tight_layout()\n",
-    "plt.savefig(\"plots/loss_curve.png\", dpi=150)\n",
+    "plt.savefig(\"plots/before_after.png\", dpi=150, bbox_inches=\"tight\")\n",
     "plt.show()\n",
-    "print(\"plots/loss_curve.png saved. Commit it to your repo now.\")"
+    "print(\"Saved: plots/before_after.png\")\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "cell 13"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "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\")"
    ]
   }
  ],

server/app.py

@@ -44,7 +44,7 @@ app = create_app(
 )
 
 
-def main(host: str = "0.0.0.0", port: int = 7068):
+def main(host: str = "0.0.0.0", port: int = 7860):
     """
     Entry point for direct execution via uv run or python -m.
 
@@ -58,6 +58,9 @@ def main(host: str = "0.0.0.0", port: int = 7068):
     """
     import uvicorn
 
+    # HuggingFace Spaces expects the app on PORT (typically 7860).
+    port = int(os.getenv("PORT", str(port)))
+
     uvicorn.run(app, host=host, port=port)