{
  "cells": [
    {
      "cell_type": "markdown",
      "id": "896b8c3f",
      "metadata": {},
      "source": [
        "# CounterFeint - Baseline Evaluation (BEFORE training)\n",
        "\n",
        "This notebook produces the **BEFORE** numbers your submission compares against.\n",
        "\n",
        "## What it does\n",
        "\n",
        "1. Loads one or more base models (no LoRA, no fine-tuning) as the Investigator\n",
        "2. Runs each model on a fixed set of held-out seeds (`task_1`, `task_2`, `task_3`)\n",
        "3. Reports `grader_score`, `verdicts`, `fallback_rate` per task per seed\n",
        "4. Saves results to `baseline_outputs/<model_tag>/baseline_results.json`\n",
        "5. Renders a single comparison plot across models\n",
        "\n",
        "## Why a separate notebook\n",
        "\n",
        "Running this **before** training gives you:\n",
        "\n",
        "- A cheap (~30 min) reproducibility check on HF Spaces hardware\n",
        "- A multi-model bake-off so you pick the best base before training\n",
        "- A reusable JSON file the training notebook can load for its `BEFORE` panel without re-running the base model\n",
        "\n",
        "## Models compared\n",
        "\n",
        "| Model | Params | Why |\n",
        "|---|---|---|\n",
        "| `Qwen/Qwen3-0.6B`             | 0.6B | Smallest Qwen3, default training target |\n",
        "| `Qwen/Qwen2.5-1.5B-Instruct`  | 1.5B | Production Investigator (current baseline) |\n",
        "| `Qwen/Qwen3-1.7B`             | 1.7B | Bigger Qwen3, fits on T4 with 4-bit |\n",
        "\n",
        "You can edit the `MODELS` list in Section 2 to add/remove.\n",
        "\n",
        "## Hardware\n",
        "\n",
        "Runs on a single T4 (~16 GB VRAM). Each model takes ~10 min for 9 episodes (3 tasks × 3 seeds)."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "7fb0c26a",
      "metadata": {},
      "source": [
        "## Section 1 - Setup\n",
        "\n",
        "Same install logic as `official_hf_training.ipynb`. If Colab asks you to restart after install, do so and resume from Section 2."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "dc628696",
      "metadata": {},
      "outputs": [],
      "source": [
        "import os\n",
        "import sys\n",
        "import subprocess\n",
        "from pathlib import Path\n",
        "\n",
        "IN_COLAB = \"google.colab\" in sys.modules\n",
        "print(f\"In Colab: {IN_COLAB}\")\n",
        "\n",
        "if IN_COLAB:\n",
        "    repo_dir = Path(\"/content/OpenEnv-Hackathon\")\n",
        "    if not repo_dir.exists():\n",
        "        subprocess.run(\n",
        "            [\"git\", \"clone\", \"https://github.com/Abhijithreddydasari/OpenEnv-Hackathon.git\", str(repo_dir)],\n",
        "            check=True,\n",
        "        )\n",
        "    REPO_ROOT = repo_dir\n",
        "else:\n",
        "    # On HF Spaces the kernel may start in /data or /home/user\n",
        "    _candidates = [\n",
        "        Path('/data/counterfeint'),\n",
        "        Path('/home/user/app/counterfeint'),\n",
        "        Path('/home/user/app'),\n",
        "    ]\n",
        "    here = Path.cwd().resolve()\n",
        "    REPO_ROOT = next(\n",
        "        (p for p in [here, *here.parents, *_candidates] if (p / 'counterfeint' / 'server').exists() or (p / 'server').exists()),\n",
        "        here,\n",
        "    )\n",
        "    # If we found a path like /data/counterfeint where server/ is directly inside,\n",
        "    # we need to go one level up for the repo root\n",
        "    if (REPO_ROOT / 'server').exists() and not (REPO_ROOT / 'counterfeint').exists():\n",
        "        REPO_ROOT = REPO_ROOT.parent\n",
        "\n",
        "print(f\"REPO_ROOT = {REPO_ROOT}\")\n",
        "os.chdir(REPO_ROOT)\n",
        "if str(REPO_ROOT) not in sys.path:\n",
        "    sys.path.insert(0, str(REPO_ROOT))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "493dc487",
      "metadata": {},
      "outputs": [],
      "source": [
        "def pip_install(args):\n",
        "    subprocess.run([sys.executable, \"-m\", \"pip\", \"install\", \"-q\", *args], check=True)\n",
        "\n",
        "print(\"Installing CounterFeint (editable) ...\")\n",
        "pip_install([\"-e\", \".\"])\n",
        "\n",
        "print(\"Installing eval dependencies ...\")\n",
        "pip_install([\n",
        "    \"torch\", \"transformers>=4.46.0\", \"accelerate>=1.1.0\",\n",
        "    \"bitsandbytes>=0.44.0\", \"datasets>=3.0.0\",\n",
        "    \"matplotlib>=3.8.0\", \"huggingface_hub>=0.26.0\",\n",
        "    \"nest_asyncio>=1.6.0\",\n",
        "])\n",
        "print(\"Done.\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "7d5fa3de",
      "metadata": {},
      "outputs": [],
      "source": [
        "try:\n",
        "    from huggingface_hub import notebook_login\n",
        "    notebook_login()\n",
        "except Exception as exc:\n",
        "    print(f\"[hf-login] skipped: {exc}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "86d9a05e",
      "metadata": {},
      "outputs": [],
      "source": [
        "import torch\n",
        "import nest_asyncio\n",
        "nest_asyncio.apply()\n",
        "\n",
        "print(f\"PyTorch: {torch.__version__}\")\n",
        "print(f\"CUDA   : {torch.cuda.is_available()}\")\n",
        "if torch.cuda.is_available():\n",
        "    print(f\"GPU    : {torch.cuda.get_device_name(0)}\")\n",
        "    print(f\"VRAM   : {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB\")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "648c9020",
      "metadata": {},
      "source": [
        "## Section 2 - Models to evaluate\n",
        "\n",
        "Edit `MODELS` to control which base models you run. Each entry is `(repo_id, short_tag)`. The short tag is used in filenames and the comparison plot.\n",
        "\n",
        "For a quick first run, comment out the bigger models. For the final submission, run all three so you have a multi-model bake-off in your README."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "b9a87a43",
      "metadata": {},
      "outputs": [],
      "source": [
        "import json\n",
        "\n",
        "MODELS = [\n",
        "    (\"Qwen/Qwen3-0.6B\",            \"qwen3-0.6b\"),\n",
        "    (\"Qwen/Qwen2.5-1.5B-Instruct\", \"qwen2.5-1.5b\"),\n",
        "    (\"Qwen/Qwen3-1.7B\",            \"qwen3-1.7b\"),\n",
        "    # Comment models out to skip them. Each adds ~10 min on T4.\n",
        "]\n",
        "\n",
        "TASKS = [\"task_1\", \"task_2\", \"task_3\"]\n",
        "EVAL_SEEDS = [101, 103, 107]   # held-out seeds\n",
        "MAX_STEPS_PER_EPISODE = 80\n",
        "\n",
        "LOAD_IN_4BIT = True\n",
        "MAX_NEW_TOKENS = 128\n",
        "TEMPERATURE = 0.3\n",
        "\n",
        "OUT_ROOT = Path(\"baseline_outputs\")\n",
        "OUT_ROOT.mkdir(parents=True, exist_ok=True)\n",
        "\n",
        "print(f\"Will evaluate {len(MODELS)} model(s) on \"\n",
        "      f\"{len(TASKS)} tasks x {len(EVAL_SEEDS)} seeds = \"\n",
        "      f\"{len(MODELS) * len(TASKS) * len(EVAL_SEEDS)} episodes\")\n",
        "for repo, tag in MODELS:\n",
        "    print(f\"  - {tag:<14}  ({repo})\")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "bca6692a",
      "metadata": {},
      "source": [
        "## Section 3 - In-process eval driver\n",
        "\n",
        "Runs one full multi-agent episode in-process (no HTTP server). Same env code as the production HTTP path.\n",
        "\n",
        "The Fraudster + Auditor are scripted (deterministic) so the only thing that varies between models is the Investigator's policy."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "4e73367c",
      "metadata": {},
      "outputs": [],
      "source": [
        "from counterfeint.scripted import HeuristicAuditor, ReactiveFraudster\n",
        "from counterfeint.server.referee import RefereeEnvironment\n",
        "\n",
        "\n",
        "def run_episode(investigator, *, task_id, seed, max_steps=80):\n",
        "    env = RefereeEnvironment()\n",
        "    env.reset_match(task_id=task_id, seed=seed)\n",
        "\n",
        "    fraudster = ReactiveFraudster(seed=42)\n",
        "    auditor = HeuristicAuditor()\n",
        "    handlers = {\n",
        "        \"fraudster_turn\":   (fraudster,    env.build_fraudster_observation,    env.step_as_fraudster),\n",
        "        \"investigator_turn\":(investigator, env.build_investigator_observation, env.step_as_investigator),\n",
        "        \"audit_phase\":      (auditor,      env.build_auditor_observation,      env.step_as_auditor),\n",
        "    }\n",
        "\n",
        "    step_idx = 0\n",
        "    n_verdicts = 0\n",
        "    n_inv_steps = 0\n",
        "    fb_at_start = getattr(investigator, \"fallback_count\", 0)\n",
        "    calls_at_start = getattr(investigator, \"call_count\", 0)\n",
        "\n",
        "    while env.phase in handlers:\n",
        "        current_phase = env.phase\n",
        "        policy, build_obs, step_fn = handlers[current_phase]\n",
        "        if current_phase != \"audit_phase\" and step_idx >= max_steps:\n",
        "            break\n",
        "        obs = build_obs().model_dump()\n",
        "        for slot in (\"last_prompt\", \"last_completion\", \"last_error\"):\n",
        "            if hasattr(policy, slot):\n",
        "                setattr(policy, slot, None)\n",
        "        try:\n",
        "            action = policy.act(obs)\n",
        "        except Exception as exc:\n",
        "            print(f\"    [policy crash] {current_phase}: {type(exc).__name__}: {exc}\")\n",
        "            break\n",
        "        try:\n",
        "            step_fn(action)\n",
        "        except Exception as exc:\n",
        "            print(f\"    [env reject] {current_phase}: {type(exc).__name__}: {exc}\")\n",
        "            break\n",
        "        if current_phase != \"audit_phase\":\n",
        "            step_idx += 1\n",
        "        if current_phase == \"investigator_turn\":\n",
        "            n_inv_steps += 1\n",
        "            if getattr(action, \"action_type\", None) in (\"verdict\", \"link_accounts\"):\n",
        "                n_verdicts += 1\n",
        "\n",
        "    state = env.state\n",
        "    fb_used = getattr(investigator, \"fallback_count\", 0) - fb_at_start\n",
        "    calls_used = getattr(investigator, \"call_count\", 0) - calls_at_start\n",
        "\n",
        "    return {\n",
        "        \"task_id\": task_id,\n",
        "        \"seed\": seed,\n",
        "        \"grader_score\": getattr(state, \"grader_score\", None),\n",
        "        \"end_reason\":   getattr(state, \"end_reason\", None),\n",
        "        \"stopped_phase\": env.phase,\n",
        "        \"n_inv_steps\": n_inv_steps,\n",
        "        \"n_verdicts\": n_verdicts,\n",
        "        \"investigator_calls\": int(calls_used),\n",
        "        \"investigator_fallback\": int(fb_used),\n",
        "    }"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "6b65bbfe",
      "metadata": {},
      "source": [
        "## Section 4 - Run baseline eval across all models\n",
        "\n",
        "Loads each model in turn (4-bit), runs all eval episodes, then unloads it before loading the next. Memory should stay under 8 GB on T4 even for the 1.7B model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "3115c1a7",
      "metadata": {},
      "outputs": [],
      "source": [
        "import gc\n",
        "import time\n",
        "from counterfeint.agents import HFInvestigator\n",
        "\n",
        "all_results = {}  # tag -> list of episode dicts\n",
        "\n",
        "for repo, tag in MODELS:\n",
        "    print(f\"\\n{'='*70}\")\n",
        "    print(f\"Loading {tag} ({repo}) ...\")\n",
        "    print(f\"{'='*70}\")\n",
        "    t0 = time.perf_counter()\n",
        "    investigator = HFInvestigator.from_pretrained(\n",
        "        repo,\n",
        "        load_in_4bit=LOAD_IN_4BIT,\n",
        "        max_new_tokens=MAX_NEW_TOKENS,\n",
        "        temperature=TEMPERATURE,\n",
        "        do_sample=True,\n",
        "        enable_thinking=False,\n",
        "    )\n",
        "    print(f\"Loaded in {time.perf_counter() - t0:.1f}s\")\n",
        "\n",
        "    rows = []\n",
        "    eval_t0 = time.perf_counter()\n",
        "    for task_id in TASKS:\n",
        "        for seed in EVAL_SEEDS:\n",
        "            r = run_episode(\n",
        "                investigator,\n",
        "                task_id=task_id, seed=seed,\n",
        "                max_steps=MAX_STEPS_PER_EPISODE,\n",
        "            )\n",
        "            rows.append(r)\n",
        "            print(\n",
        "                f\"  {task_id} seed={seed:>3}: \"\n",
        "                f\"grader={(r['grader_score'] or 0.0):+.3f}  \"\n",
        "                f\"verdicts={r['n_verdicts']:>2}  \"\n",
        "                f\"fallback={r['investigator_fallback']}/{r['investigator_calls']}  \"\n",
        "                f\"end={r['end_reason']}\"\n",
        "            )\n",
        "    eval_elapsed = time.perf_counter() - eval_t0\n",
        "    print(f\"\\nEval took {eval_elapsed/60:.1f} min ({eval_elapsed/len(rows):.1f}s/episode)\")\n",
        "\n",
        "    out_dir = OUT_ROOT / tag\n",
        "    out_dir.mkdir(parents=True, exist_ok=True)\n",
        "    with open(out_dir / \"baseline_results.json\", \"w\") as f:\n",
        "        json.dump({\n",
        "            \"model_repo\": repo,\n",
        "            \"model_tag\":  tag,\n",
        "            \"config\": {\n",
        "                \"load_in_4bit\": LOAD_IN_4BIT,\n",
        "                \"max_new_tokens\": MAX_NEW_TOKENS,\n",
        "                \"temperature\": TEMPERATURE,\n",
        "                \"max_steps_per_episode\": MAX_STEPS_PER_EPISODE,\n",
        "                \"eval_seeds\": EVAL_SEEDS,\n",
        "                \"tasks\": TASKS,\n",
        "            },\n",
        "            \"rows\": rows,\n",
        "            \"elapsed_seconds\": round(eval_elapsed, 1),\n",
        "        }, f, indent=2)\n",
        "    print(f\"Wrote {out_dir/'baseline_results.json'}\")\n",
        "\n",
        "    all_results[tag] = rows\n",
        "\n",
        "    del investigator\n",
        "    gc.collect()\n",
        "    if torch.cuda.is_available():\n",
        "        torch.cuda.empty_cache()\n",
        "\n",
        "print(\"\\n\" + \"=\" * 70)\n",
        "print(\"ALL MODELS DONE\")\n",
        "print(\"=\" * 70)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "e149ad5a",
      "metadata": {},
      "source": [
        "## Section 5 - Summary table + comparison plot\n",
        "\n",
        "A clean per-model summary, plus the figure your README and slide deck will reference."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "4cf2ea68",
      "metadata": {},
      "outputs": [],
      "source": [
        "def _gs(r): return r[\"grader_score\"] or 0.0\n",
        "\n",
        "summary_rows = []\n",
        "for tag, rows in all_results.items():\n",
        "    by_task = {t: [] for t in TASKS}\n",
        "    for r in rows:\n",
        "        by_task[r[\"task_id\"]].append(_gs(r))\n",
        "    overall = sum(_gs(r) for r in rows) / max(1, len(rows))\n",
        "    fb = sum(r[\"investigator_fallback\"] for r in rows)\n",
        "    calls = sum(r[\"investigator_calls\"] for r in rows)\n",
        "    summary_rows.append({\n",
        "        \"model\": tag,\n",
        "        \"task_1_mean\": sum(by_task[\"task_1\"]) / max(1, len(by_task[\"task_1\"])),\n",
        "        \"task_2_mean\": sum(by_task[\"task_2\"]) / max(1, len(by_task[\"task_2\"])),\n",
        "        \"task_3_mean\": sum(by_task[\"task_3\"]) / max(1, len(by_task[\"task_3\"])),\n",
        "        \"overall_mean\": overall,\n",
        "        \"fallback_rate\": fb / max(1, calls),\n",
        "    })\n",
        "\n",
        "print(f\"\\n{'model':<14} {'task_1':>8} {'task_2':>8} {'task_3':>8} {'mean':>8} {'fb-rate':>8}\")\n",
        "print(\"-\" * 60)\n",
        "for s in summary_rows:\n",
        "    print(\n",
        "        f\"{s['model']:<14} \"\n",
        "        f\"{s['task_1_mean']:>8.3f} \"\n",
        "        f\"{s['task_2_mean']:>8.3f} \"\n",
        "        f\"{s['task_3_mean']:>8.3f} \"\n",
        "        f\"{s['overall_mean']:>8.3f} \"\n",
        "        f\"{s['fallback_rate']:>8.2%}\"\n",
        "    )\n",
        "\n",
        "with open(OUT_ROOT / \"baseline_summary.json\", \"w\") as f:\n",
        "    json.dump(summary_rows, f, indent=2)\n",
        "print(f\"\\nWrote {OUT_ROOT/'baseline_summary.json'}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "995e389c",
      "metadata": {},
      "outputs": [],
      "source": [
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "\n",
        "if summary_rows:\n",
        "    tags = [s[\"model\"] for s in summary_rows]\n",
        "    width = 0.18\n",
        "    x = np.arange(len(tags))\n",
        "\n",
        "    fig, ax = plt.subplots(figsize=(10, 5.5))\n",
        "    ax.bar(x - 1.5 * width, [s[\"task_1_mean\"] for s in summary_rows], width, label=\"task_1\", color=\"#4C72B0\")\n",
        "    ax.bar(x - 0.5 * width, [s[\"task_2_mean\"] for s in summary_rows], width, label=\"task_2\", color=\"#55A868\")\n",
        "    ax.bar(x + 0.5 * width, [s[\"task_3_mean\"] for s in summary_rows], width, label=\"task_3\", color=\"#C44E52\")\n",
        "    ax.bar(x + 1.5 * width, [s[\"overall_mean\"] for s in summary_rows], width, label=\"overall\", color=\"#8172B2\")\n",
        "\n",
        "    ax.set_xticks(x)\n",
        "    ax.set_xticklabels(tags, rotation=15, ha=\"right\")\n",
        "    ax.set_ylabel(\"Mean grader_score\")\n",
        "    ax.set_title(\"Baseline grader_score across models (no training)\")\n",
        "    ax.set_ylim(0, 1.0)\n",
        "    ax.axhline(0.45, ls=\"--\", color=\"gray\", lw=1, label=\"trainable bar\")\n",
        "    ax.legend(ncol=5, fontsize=9, loc=\"upper right\")\n",
        "    ax.grid(axis=\"y\", alpha=0.3)\n",
        "    fig.tight_layout()\n",
        "    plot_path = OUT_ROOT / \"baseline_comparison.png\"\n",
        "    fig.savefig(plot_path, dpi=140, bbox_inches=\"tight\")\n",
        "    plt.show()\n",
        "    print(f\"\\nWrote {plot_path}\")\n",
        "else:\n",
        "    print(\"No results to plot.\")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "802d3ec2",
      "metadata": {},
      "source": [
        "## Section 6 - (optional) Push artifacts to HF Hub\n",
        "\n",
        "Pushes the baseline JSON + plot to a HF dataset repo so judges can browse them. Set `BASELINE_HUB_REPO_ID` and re-run."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "dcefbda3",
      "metadata": {},
      "outputs": [],
      "source": [
        "BASELINE_HUB_REPO_ID = \"\"    # e.g. \"your-username/counterfeint-baselines\"\n",
        "\n",
        "if BASELINE_HUB_REPO_ID:\n",
        "    from huggingface_hub import HfApi, create_repo\n",
        "    api = HfApi()\n",
        "    create_repo(BASELINE_HUB_REPO_ID, repo_type=\"dataset\", exist_ok=True, private=False)\n",
        "    api.upload_folder(\n",
        "        folder_path=str(OUT_ROOT),\n",
        "        repo_id=BASELINE_HUB_REPO_ID,\n",
        "        repo_type=\"dataset\",\n",
        "        commit_message=\"Upload CounterFeint baseline eval artifacts\",\n",
        "    )\n",
        "    print(f\"Pushed to https://huggingface.co/datasets/{BASELINE_HUB_REPO_ID}\")\n",
        "else:\n",
        "    print(\"Set BASELINE_HUB_REPO_ID to push results.\")"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "name": "python",
      "version": "3.11"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}