{ "nbformat": 4, "nbformat_minor": 5, "metadata": { "kernelspec": { "display_name": "Python 3", "language": "python", "name": "python3" }, "language_info": { "name": "python", "version": "3.10.0" } }, "cells": [ { "cell_type": "markdown", "id": "title", "metadata": {}, "source": [ "# AP Commander — Baseline Eval + GRPO Training\n", "\n", "**Run on:** Google Colab (T4/A100) · HF Spaces (A10G) · any CUDA machine \n", "**Environment:** `https://pathikreet-ap-clerk-env.hf.space` (always live) \n", "**Repo:** `Pathikreet/ap-commander-training`\n", "\n", "```\n", "This notebook has two independent sections:\n", " Part A — Baseline Eval (no training, measure untrained model)\n", " Part B — GRPO Training (fine-tune with curriculum + accumulated rewards)\n", "```\n", "\n", "Run Part A first to get the baseline, then Part B to train. Each run saves to a\n", "timestamped folder under `runs/` so nothing is ever overwritten." ] }, { "cell_type": "code", "execution_count": null, "id": "check_gpu", "metadata": {}, "outputs": [], "source": [ "import torch\n", "assert torch.cuda.is_available(), 'No GPU — go to Runtime > Change runtime type > GPU (T4 or A100)'\n", "gpu = torch.cuda.get_device_properties(0)\n", "vram = gpu.total_memory / 1e9\n", "print(f'GPU : {gpu.name}')\n", "print(f'VRAM: {vram:.1f} GB')\n", "if vram < 14:\n", " print('WARNING: < 14 GB VRAM — use Qwen2.5-1.5B or reduce batch size')\n", "else:\n", " print('OK: enough VRAM for Qwen2.5-7B in 4-bit')" ] }, { "cell_type": "code", "execution_count": null, "id": "install", "metadata": {}, "outputs": [], "source": [ "# Install — TRL standard stack (no Unsloth: avoids llm_blender conflicts on HF Spaces)\n", "!pip install -q 'trl>=0.15.0' accelerate peft transformers bitsandbytes\n", "!pip install -q requests datasets matplotlib huggingface_hub\n", "print('Done')" ] }, { "cell_type": "code", "execution_count": null, "id": "config", "metadata": {}, "outputs": [], "source": [ "import os\n", "\n", "# ── CONFIG — edit these ──────────────────────────────────────────────────────\n", "MODEL_NAME = 'Qwen/Qwen2.5-7B-Instruct' # swap to 1.5B for quick test\n", "# MODEL_NAME = 'Qwen/Qwen2.5-1.5B-Instruct'\n", "# MODEL_NAME = 'meta-llama/Meta-Llama-3-8B-Instruct' # needs HF_TOKEN\n", "\n", "HF_TOKEN = '' # paste hf_... here, or set as Colab secret\n", "NUM_EPOCHS = 3\n", "NUM_GENERATIONS = 8 # GRPO group size; set 4 on T4 to save memory\n", "EVAL_SEEDS = [42, 99, 7] # 3 seeds per task for baseline\n", "ENV_URL = 'https://pathikreet-ap-clerk-env.hf.space'\n", "# ─────────────────────────────────────────────────────────────────────────────\n", "\n", "# Auto-load from env if running on HF Spaces\n", "if not HF_TOKEN:\n", " HF_TOKEN = os.environ.get('HF_TOKEN', '')\n", "\n", "if HF_TOKEN:\n", " from huggingface_hub import login\n", " login(token=HF_TOKEN, add_to_git_credential=False)\n", " print('[AUTH] Logged in to HF Hub')\n", "else:\n", " print('[AUTH] No HF_TOKEN — public models only (Qwen recommended)')\n", "\n", "print(f'Model : {MODEL_NAME}')\n", "print(f'Epochs: {NUM_EPOCHS} | Generations: {NUM_GENERATIONS}')" ] }, { "cell_type": "code", "execution_count": null, "id": "env_check", "metadata": {}, "outputs": [], "source": [ "import requests\n", "\n", "h = requests.get(f'{ENV_URL}/health', timeout=30).json()\n", "print(f\"Environment: {h['status']} | tasks: {h.get('total_tasks')} | version: {h.get('version')}\")\n", "\n", "# Quick sanity: one episode\n", "reset = requests.post(f'{ENV_URL}/reset', json={'task_id': 'easy_perfect_match', 'seed': 42}).json()\n", "step = requests.post(f'{ENV_URL}/step', json={\n", " 'session_id': reset['session_id'],\n", " 'action': {'decision': 'APPROVE_FULL',\n", " 'approved_amount': reset['observation']['invoice']['invoice_total'],\n", " 'reason_code': 'MATCH_CONFIRMED',\n", " 'explanation': 'Invoice matches PO and GRN. Three-way match confirmed.'}\n", "}).json()\n", "print(f\"Sanity check: score={step['reward']['score']} (expect ~0.99)\")" ] }, { "cell_type": "code", "execution_count": null, "id": "helpers", "metadata": {}, "outputs": [], "source": [ "import json, re, random, time, datetime, collections, os\n", "import matplotlib\n", "matplotlib.use('Agg')\n", "import matplotlib.pyplot as plt\n", "import numpy as np\n", "\n", "SYSTEM_PROMPT = \"\"\"You are an AI Accounts Payable Clerk. Review the invoice, PO, and GRN, then output ONLY valid JSON:\n", "{\"decision\": \"APPROVE_FULL\"|\"APPROVE_PARTIAL\"|\"REJECT\"|\"ESCALATE\"|\"QUERY_VENDOR\",\n", " \"approved_amount\": ,\n", " \"reason_code\": \"MATCH_CONFIRMED\"|\"QUANTITY_MISMATCH\"|\"PRICE_DISCREPANCY\"|\"POLICY_VIOLATION\"|\"NO_PO_FOUND\"|\"DUPLICATE_INVOICE\"|\"VENDOR_MISMATCH\"|\"TAX_DISCREPANCY\"|\"PENDING_CLARIFICATION\"|\"MANAGER_REVIEW\",\n", " \"explanation\": \"\"}\"\"\"\n", "\n", "VALID_DECISIONS = {'APPROVE_FULL','APPROVE_PARTIAL','REJECT','ESCALATE','QUERY_VENDOR','HOLD'}\n", "VALID_REASON_CODES = {'MATCH_CONFIRMED','QUANTITY_MISMATCH','PRICE_DISCREPANCY','POLICY_VIOLATION',\n", " 'NO_PO_FOUND','DUPLICATE_INVOICE','VENDOR_MISMATCH','TAX_DISCREPANCY',\n", " 'PENDING_CLARIFICATION','MANAGER_REVIEW'}\n", "\n", "EVAL_TASKS = [\n", " 'easy_perfect_match', 'easy_no_po_found',\n", " 'medium_quantity_shortfall', 'medium_price_discrepancy',\n", " 'medium_split_delivery', 'medium_vendor_mismatch',\n", " 'hard_policy_violation', 'hard_duplicate_invoice',\n", " 'hard_partial_po_match', 'hard_tax_discrepancy',\n", "]\n", "TRAIN_TASKS = EVAL_TASKS\n", "\n", "TASK_DIFFICULTY = {\n", " 'easy_perfect_match': 'easy', 'easy_no_po_found': 'easy',\n", " 'medium_quantity_shortfall':'medium','medium_price_discrepancy':'medium',\n", " 'medium_split_delivery':'medium', 'medium_vendor_mismatch':'medium',\n", " 'hard_policy_violation':'hard', 'hard_duplicate_invoice':'hard',\n", " 'hard_partial_po_match':'hard', 'hard_tax_discrepancy':'hard',\n", "}\n", "DIFF_COLORS = {'easy': '#3fb950', 'medium': '#d29922', 'hard': '#f85149'}\n", "\n", "\n", "def obs_to_prompt(obs):\n", " inv = obs['invoice']\n", " lines = '\\n'.join(f\" {li['description']}: qty={li['quantity']}, unit_price=${li['unit_price']:.2f}\"\n", " for li in inv.get('line_items', []))\n", " pos = '\\n'.join(\n", " f\" PO {p['po_number']} ({p['status']}) {p['vendor_name']}: \" +\n", " ', '.join(f\"{l['description']} qty={l['ordered_quantity']} @${l['agreed_unit_price']:.2f}\"\n", " for l in p.get('lines', []))\n", " for p in obs.get('purchase_orders', []))\n", " grns = '\\n'.join(\n", " f\" GRN {g['grn_id']}: \" + ', '.join(f\"{l['description']} recv={l['received_quantity']}\"\n", " for l in g.get('lines', []))\n", " for g in obs.get('goods_receipts', []))\n", " ctx = '\\n'.join(f' {n}' for n in obs.get('context_notes', []))\n", " paid = ', '.join(obs.get('paid_invoice_ids', []))\n", " return (\n", " f\"TASK: {obs['task_name']}\\n{obs['task_description']}\\n\\n\"\n", " f\"INVOICE {inv['invoice_id']} | {inv['vendor_name']} | ${inv['invoice_total']:,.2f}\\n{lines}\\n\"\n", " f\"Freight: ${inv.get('freight_charge',0):.2f}\\n\\n\"\n", " f\"PURCHASE ORDERS:\\n{pos}\\n\\nGOODS RECEIPTS:\\n{grns}\\n\"\n", " + (f\"PAID LEDGER: {paid}\\n\" if paid else \"\")\n", " + (f\"CONTEXT:\\n{ctx}\\n\" if ctx else \"\")\n", " + f\"\\nPOLICY:\\n{obs['company_policy']}\\n\\nOutput JSON decision.\"\n", " )\n", "\n", "\n", "def parse_action(raw):\n", " clean = re.sub(r'```(?:json)?\\s*|\\s*```', '', raw).strip()\n", " m = re.search(r'\\{.*\\}', clean, re.DOTALL)\n", " if m:\n", " try:\n", " a = json.loads(m.group())\n", " if (a.get('decision') in VALID_DECISIONS and\n", " a.get('reason_code') in VALID_REASON_CODES and\n", " isinstance(a.get('approved_amount'), (int, float)) and\n", " len(a.get('explanation', '')) > 10):\n", " return a, True\n", " except Exception:\n", " pass\n", " return {'decision': 'REJECT', 'approved_amount': 0.0,\n", " 'reason_code': 'NO_PO_FOUND', 'explanation': 'parse error'}, False\n", "\n", "\n", "def _greedy_followup(obs_dict):\n", " notes = ' '.join(obs_dict.get('context_notes', [])).lower()\n", " total = abs(float(obs_dict.get('invoice', {}).get('invoice_total', 0) or 0))\n", " if any(k in notes for k in ('manager approved', 'vp approved', 'cfo approved', 'approved by')):\n", " return {'decision': 'APPROVE_FULL', 'approved_amount': total,\n", " 'reason_code': 'MATCH_CONFIRMED', 'explanation': f'Approval confirmed. Approving ${total:.2f}.'}\n", " if any(k in notes for k in ('fraudulent', 'duplicate', 'already paid', 'deny')):\n", " return {'decision': 'REJECT', 'approved_amount': 0.0,\n", " 'reason_code': 'DUPLICATE_INVOICE', 'explanation': 'Confirmed fraud/duplicate. Rejecting.'}\n", " if any(k in notes for k in ('flagged', 'violation', 'sox', 'non-compliant')):\n", " return {'decision': 'REJECT', 'approved_amount': 0.0,\n", " 'reason_code': 'POLICY_VIOLATION', 'explanation': 'Compliance violation. Rejecting.'}\n", " return {'decision': 'REJECT', 'approved_amount': 0.0,\n", " 'reason_code': 'PENDING_CLARIFICATION', 'explanation': 'Could not resolve. Rejecting for safety.'}\n", "\n", "\n", "def run_episode_accumulated(task_id, first_action, seed=None, discount=0.9, max_steps=5):\n", " \"\"\"Full multi-step rollout with discounted reward accumulation.\n", " QUERY_VENDOR->REJECT = 0.01 + 0.9*0.99 = 0.901 > shortcut REJECT ~0.4\"\"\"\n", " try:\n", " r = requests.post(f'{ENV_URL}/reset', json={'task_id': task_id, 'seed': seed}, timeout=20)\n", " session_id = r.json()['session_id']\n", " action, total = first_action, 0.0\n", " for step_n in range(max_steps):\n", " result = requests.post(f'{ENV_URL}/step',\n", " json={'session_id': session_id, 'action': action},\n", " timeout=20).json()\n", " total += (discount ** step_n) * float(result['reward']['score'])\n", " if result['done']:\n", " break\n", " action = _greedy_followup(result['observation'])\n", " return min(0.99, max(0.01, total))\n", " except Exception:\n", " return 0.01\n", "\n", "\n", "# ── Dark-theme plot helpers ───────────────────────────────────────────────────\n", "BG, PANEL, GRID = '#0d1117', '#161b22', '#21262d'\n", "TXT, DIM = '#e6edf3', '#8b949e'\n", "\n", "def dark_fig(*args, **kwargs):\n", " fig = plt.figure(*args, **kwargs)\n", " fig.patch.set_facecolor(BG)\n", " return fig\n", "\n", "def style_ax(ax, title='', xlabel='', ylabel=''):\n", " ax.set_facecolor(PANEL)\n", " ax.tick_params(colors='#c9d1d9', labelsize=8)\n", " for sp in ax.spines.values(): sp.set_color('#30363d')\n", " ax.spines['top'].set_visible(False); ax.spines['right'].set_visible(False)\n", " ax.yaxis.grid(True, color=GRID, linewidth=0.7); ax.set_axisbelow(True)\n", " if title: ax.set_title(title, color=TXT, fontsize=10, fontweight='bold', pad=8)\n", " if xlabel: ax.set_xlabel(xlabel, color=DIM, fontsize=8)\n", " if ylabel: ax.set_ylabel(ylabel, color=DIM, fontsize=8)\n", "\n", "print('Helper functions ready.')" ] }, { "cell_type": "markdown", "id": "part_a_header", "metadata": {}, "source": [ "---\n", "## Part A — Untrained Baseline Evaluation\n", "Loads the model **without any fine-tuning**, evaluates all 10 tasks × 3 seeds,\n", "produces plots, and saves to `runs/baselines/MODEL-DATETIME/`.\n", "\n", "**Skip to Part B** if you only want to train." ] }, { "cell_type": "code", "execution_count": null, "id": "load_model_base", "metadata": {}, "outputs": [], "source": [ "import torch\n", "from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig\n", "\n", "bnb = BitsAndBytesConfig(\n", " load_in_4bit=True, bnb_4bit_quant_type='nf4',\n", " bnb_4bit_compute_dtype=torch.bfloat16, bnb_4bit_use_double_quant=True,\n", ")\n", "tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)\n", "if tokenizer.pad_token is None:\n", " tokenizer.pad_token = tokenizer.eos_token\n", "\n", "model_base = AutoModelForCausalLM.from_pretrained(\n", " MODEL_NAME, quantization_config=bnb, device_map='auto', trust_remote_code=True,\n", ")\n", "model_base.eval()\n", "print(f'Loaded {MODEL_NAME} (4-bit NF4, no LoRA)')" ] }, { "cell_type": "code", "execution_count": null, "id": "run_baseline", "metadata": {}, "outputs": [], "source": [ "def eval_one(model, task_id, seed):\n", " try:\n", " reset = requests.post(f'{ENV_URL}/reset', json={'task_id': task_id, 'seed': seed}, timeout=20).json()\n", " obs, sid = reset['observation'], reset['session_id']\n", " msgs = [{'role':'system','content':SYSTEM_PROMPT},\n", " {'role':'user','content':obs_to_prompt(obs)}]\n", " text = tokenizer.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True)\n", " inputs = tokenizer(text, return_tensors='pt').to('cuda')\n", " with torch.no_grad():\n", " out = model.generate(**inputs, max_new_tokens=250, temperature=0.1, do_sample=True)\n", " raw = tokenizer.decode(out[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True)\n", " action, fmt_ok = parse_action(raw)\n", " score = float(requests.post(f'{ENV_URL}/step',\n", " json={'session_id': sid, 'action': action},\n", " timeout=20).json()['reward']['score'])\n", " return score, raw[:120], action.get('decision','?'), fmt_ok\n", " except Exception as e:\n", " return 0.01, str(e), 'ERROR', False\n", "\n", "\n", "print(f'Evaluating {len(EVAL_TASKS)} tasks × {len(EVAL_SEEDS)} seeds = {len(EVAL_TASKS)*len(EVAL_SEEDS)} episodes\\n')\n", "baseline_results = {}\n", "parse_failures = 0\n", "\n", "for task_id in EVAL_TASKS:\n", " diff = TASK_DIFFICULTY[task_id]\n", " scores, decisions, fmts = [], [], []\n", " for seed in EVAL_SEEDS:\n", " score, raw, dec, fmt_ok = eval_one(model_base, task_id, seed)\n", " scores.append(score); decisions.append(dec); fmts.append(fmt_ok)\n", " if not fmt_ok: parse_failures += 1\n", " print(f' [{diff[:4]}] {task_id:<35} seed={seed} score={score:.3f} {dec} fmt={fmt_ok}')\n", " print(f' {raw[:90]}')\n", " time.sleep(0.2)\n", " baseline_results[task_id] = {\n", " 'difficulty': diff,\n", " 'scores': [round(s,4) for s in scores],\n", " 'mean': round(sum(scores)/len(scores), 4),\n", " 'decisions': decisions,\n", " 'fmt_rate': round(sum(fmts)/len(fmts), 3),\n", " }\n", "\n", "by_diff = {}\n", "for tid, v in baseline_results.items():\n", " by_diff.setdefault(v['difficulty'], []).append(v['mean'])\n", "\n", "overall_baseline = sum(v['mean'] for v in baseline_results.values()) / len(baseline_results)\n", "print(f'\\nOverall mean: {overall_baseline:.3f} parse_failures: {parse_failures}/{len(EVAL_TASKS)*len(EVAL_SEEDS)}')\n", "for d in ['easy','medium','hard']:\n", " ms = by_diff.get(d,[])\n", " if ms: print(f' {d}: {sum(ms)/len(ms):.3f}')" ] }, { "cell_type": "code", "execution_count": null, "id": "plot_baseline", "metadata": {}, "outputs": [], "source": [ "fig = dark_fig(figsize=(16, 6))\n", "gs = fig.add_gridspec(1, 2, wspace=0.32)\n", "\n", "# Panel 1 — per-task horizontal bars\n", "ax1 = fig.add_subplot(gs[0,0])\n", "tasks = list(baseline_results.keys())\n", "means = [baseline_results[t]['mean'] for t in tasks]\n", "colors = [DIFF_COLORS[baseline_results[t]['difficulty']] for t in tasks]\n", "short = [t.replace('easy_','').replace('medium_','').replace('hard_','').replace('_',' ').title()\n", " for t in tasks]\n", "yp = range(len(tasks))\n", "ax1.barh(list(yp), means, color=colors, alpha=0.85, edgecolor=BG)\n", "ax1.set_yticks(list(yp)); ax1.set_yticklabels(short, fontsize=8)\n", "ax1.set_xlim(0, 1.05)\n", "ax1.axvline(overall_baseline, color='#f78166', linestyle='--', linewidth=1.2,\n", " label=f'Mean {overall_baseline:.3f}')\n", "ax1.axvline(0.5, color='#484f58', linestyle=':', linewidth=1)\n", "for i,m in enumerate(means):\n", " ax1.text(m+0.01, i, f'{m:.3f}', va='center', color='#c9d1d9', fontsize=8)\n", "from matplotlib.patches import Patch\n", "leg = [Patch(facecolor=c, label=d) for d,c in DIFF_COLORS.items()]\n", "leg.append(plt.Line2D([0],[0], color='#f78166', linestyle='--', label=f'Mean {overall_baseline:.3f}'))\n", "ax1.legend(handles=leg, fontsize=8, facecolor=PANEL, edgecolor='#30363d',\n", " labelcolor='#c9d1d9', loc='lower right')\n", "style_ax(ax1, f'Per-Task Baseline ({len(EVAL_SEEDS)} seeds)', ylabel='Score')\n", "\n", "# Panel 2 — mean by difficulty\n", "ax2 = fig.add_subplot(gs[0,1])\n", "diffs = ['easy','medium','hard']\n", "dmeans = [sum(by_diff.get(d,[0]))/max(1,len(by_diff.get(d,[0]))) for d in diffs]\n", "bars = ax2.bar(diffs, dmeans, color=[DIFF_COLORS[d] for d in diffs],\n", " alpha=0.85, edgecolor=BG, width=0.5)\n", "for i,(d,m) in enumerate(zip(diffs,dmeans)):\n", " ax2.text(i, m+0.02, f'{m:.3f}', ha='center', color='#c9d1d9', fontsize=11, fontweight='bold')\n", "ax2.set_ylim(0, 1.05)\n", "ax2.axhline(overall_baseline, color='#f78166', linestyle='--', linewidth=1,\n", " label=f'Overall {overall_baseline:.3f}')\n", "ax2.legend(fontsize=8, facecolor=PANEL, edgecolor='#30363d', labelcolor='#c9d1d9')\n", "style_ax(ax2, 'Mean by Difficulty', ylabel='Mean Score')\n", "\n", "model_short = MODEL_NAME.split('/')[-1]\n", "fig.suptitle(f'{model_short} — Untrained Baseline | 4-bit NF4 | {len(EVAL_SEEDS)} seeds | '\n", " f'overall={overall_baseline:.3f} | {datetime.datetime.now().strftime(\"%Y-%m-%d\")}',\n", " color=TXT, fontsize=10, y=1.01)\n", "plt.tight_layout()\n", "plt.savefig('/tmp/baseline_plot.png', dpi=130, bbox_inches='tight', facecolor=BG)\n", "plt.show()\n", "print('Plot saved: /tmp/baseline_plot.png')" ] }, { "cell_type": "code", "execution_count": null, "id": "save_baseline", "metadata": {}, "outputs": [], "source": [ "import shutil\n", "\n", "model_slug = MODEL_NAME.split('/')[-1].lower().replace('.', '-')\n", "ts = datetime.datetime.now().strftime('%Y-%m-%d_%H%M')\n", "run_dir = f'/tmp/runs/baselines/{model_slug}-{ts}'\n", "os.makedirs(run_dir, exist_ok=True)\n", "\n", "output = {\n", " 'run_type': 'llm_baseline_no_finetuning',\n", " 'model': MODEL_NAME,\n", " 'quantization': '4-bit NF4',\n", " 'lora': None,\n", " 'timestamp': datetime.datetime.now().isoformat(),\n", " 'env_url': ENV_URL,\n", " 'seeds': EVAL_SEEDS,\n", " 'eval_tasks': EVAL_TASKS,\n", " 'overall_mean': round(overall_baseline, 4),\n", " 'parse_failures': parse_failures,\n", " 'by_difficulty': {d: round(sum(ms)/len(ms),4) for d,ms in by_diff.items()},\n", " 'tasks': baseline_results,\n", "}\n", "json_path = os.path.join(run_dir, 'baseline_results.json')\n", "with open(json_path, 'w') as f:\n", " json.dump(output, f, indent=2)\n", "shutil.copy('/tmp/baseline_plot.png', os.path.join(run_dir, 'baseline_plot.png'))\n", "\n", "print(f'Saved JSON and plot to {run_dir}')\n", "\n", "# Upload to HF Space repo (runs/baselines/MODEL-DATETIME/)\n", "if HF_TOKEN:\n", " try:\n", " from huggingface_hub import HfApi\n", " api = HfApi(token=HF_TOKEN)\n", " api.upload_folder(\n", " folder_path=run_dir,\n", " path_in_repo=f'runs/baselines/{model_slug}-{ts}',\n", " repo_id='Pathikreet/ap-commander-training',\n", " repo_type='space',\n", " commit_message=f'Baseline: {model_slug} untrained {ts}',\n", " )\n", " print(f'Uploaded to HF Space repo: runs/baselines/{model_slug}-{ts}/')\n", " except Exception as e:\n", " print(f'Upload skipped: {e}')\n", "else:\n", " print('No HF_TOKEN — skipping upload. Files saved locally to', run_dir)" ] }, { "cell_type": "markdown", "id": "part_b_header", "metadata": {}, "source": [ "---\n", "## Part B — GRPO Training\n", "\n", "Adds LoRA to the loaded model (or reloads fresh), builds a curriculum-weighted\n", "dataset, trains with two independent reward functions, and produces the full\n", "4-panel results figure. Run Part A first to get `model_base` and `baseline_results`.\n", "\n", "| Setting | Value | Why |\n", "|---|---|---|\n", "| `per_device_train_batch_size = num_generations` | 8 | TRL divisibility requirement |\n", "| Two reward functions | env + format | Separate signals per guide |\n", "| `run_episode_accumulated()` | discount=0.9 | QUERY_VENDOR→REJECT = 0.901 > shortcut REJECT 0.4 |\n", "| `CurriculumSampler` | easy→medium→hard | Unlocks harder tasks as easy mean ≥ 0.70 |" ] }, { "cell_type": "code", "execution_count": null, "id": "add_lora", "metadata": {}, "outputs": [], "source": [ "from peft import LoraConfig, get_peft_model, TaskType\n", "\n", "# Add LoRA on top of the quantized base model\n", "# (model_base was loaded in Part A; if skipping Part A, re-run load_model_base first)\n", "model_base.train()\n", "model_base.enable_input_require_grads()\n", "model_base.gradient_checkpointing_enable()\n", "\n", "lora_cfg = LoraConfig(\n", " r=16, lora_alpha=16,\n", " target_modules=['q_proj','k_proj','v_proj','o_proj','gate_proj','up_proj','down_proj'],\n", " lora_dropout=0, bias='none',\n", " task_type=TaskType.CAUSAL_LM,\n", ")\n", "model = get_peft_model(model_base, lora_cfg)\n", "model.print_trainable_parameters()" ] }, { "cell_type": "code", "execution_count": null, "id": "curriculum", "metadata": {}, "outputs": [], "source": [ "_DIFFICULTY_ORDER = ['easy', 'medium', 'hard']\n", "_UNLOCK_THRESHOLDS = {'easy': 0.70, 'medium': 0.65}\n", "\n", "class CurriculumSampler:\n", " def __init__(self):\n", " self._rewards = collections.defaultdict(list)\n", " self.unlocked = {'easy'}\n", " def record(self, tid, r):\n", " self._rewards[tid].append(r)\n", " self._try_unlock()\n", " def mean_for(self, diff):\n", " v = [r for tid,d in TASK_DIFFICULTY.items() if d==diff for r in self._rewards.get(tid,[])]\n", " return sum(v)/len(v) if v else 0.0\n", " def _try_unlock(self):\n", " for i, d in enumerate(_DIFFICULTY_ORDER[:-1]):\n", " if d in self.unlocked and self.mean_for(d) >= _UNLOCK_THRESHOLDS.get(d, 0.70):\n", " nxt = _DIFFICULTY_ORDER[i+1]\n", " if nxt not in self.unlocked:\n", " self.unlocked.add(nxt)\n", " print(f'[CURRICULUM] Unlocked {nxt}! mean({d})={self.mean_for(d):.3f}')\n", " def gate_task(self, tid):\n", " if TASK_DIFFICULTY.get(tid,'easy') in self.unlocked:\n", " return tid\n", " return random.choice([t for t,d in TASK_DIFFICULTY.items() if d=='easy'])\n", " def build_dataset_tasks(self):\n", " seeds_per_diff = {'easy':10, 'medium':5, 'hard':2}\n", " return [(tid,s) for tid,diff in TASK_DIFFICULTY.items()\n", " if diff in self.unlocked\n", " for s in range(1, seeds_per_diff[diff]+1)]\n", " def status(self):\n", " return ' | '.join(f'{d}={self.mean_for(d):.2f}{\"✓\" if d in self.unlocked else \"✗\"}'\n", " for d in _DIFFICULTY_ORDER)\n", "\n", "CURRICULUM = CurriculumSampler()\n", "\n", "\n", "# Metrics tracker for live plots\n", "class Metrics:\n", " def __init__(self):\n", " self.step = 0\n", " self.reward_history = [] # [(step, mean_reward)]\n", " self.decision_counts= collections.Counter()\n", " self.format_scores = []\n", " self.reward_by_task = collections.defaultdict(list)\n", " self.parse_failures = 0\n", " self.total_calls = 0\n", " self._t0 = time.time()\n", " def log(self, rewards, decisions, fmt_oks, task_ids):\n", " self.step += 1\n", " self.total_calls+= len(rewards)\n", " self.reward_history.append((self.step, sum(rewards)/len(rewards)))\n", " for d in decisions: self.decision_counts[d] += 1\n", " for ok in fmt_oks: self.format_scores.append(1.0 if ok else 0.0)\n", " for tid,r in zip(task_ids,rewards): self.reward_by_task[tid].append(r)\n", " def recent_mean(self, n=20):\n", " tail = self.reward_history[-n:]\n", " return sum(r for _,r in tail)/len(tail) if tail else 0.0\n", " def fmt_rate(self):\n", " return sum(self.format_scores)/len(self.format_scores) if self.format_scores else 0.0\n", " def summary(self):\n", " task_means = {t:round(sum(v)/len(v),3) for t,v in self.reward_by_task.items()}\n", " print(f'[METRICS] step={self.step} recent={self.recent_mean():.3f} '\n", " f'fmt={self.fmt_rate():.1%} parse_fails={self.parse_failures} '\n", " f'total_calls={self.total_calls}')\n", " print(f'[METRICS] per_task: {task_means}')\n", " print(f'[CURRICULUM] {CURRICULUM.status()}')\n", "\n", "METRICS = Metrics()\n", "print('Curriculum and Metrics ready. Curriculum unlocked:', CURRICULUM.unlocked)" ] }, { "cell_type": "code", "execution_count": null, "id": "reward_fns", "metadata": {}, "outputs": [], "source": [ "LOG_EVERY = 20\n\ndef env_reward_fn(completions, task_id=None, seed=None, **kwargs):\n \"\"\"Accumulated discounted per-step reward. Curriculum gating active.\"\"\"\n task_ids = task_id if task_id else ['easy_perfect_match']*len(completions)\n seeds = seed if seed else [random.randint(1,999)]*len(completions)\n rewards, decisions, fmts = [], [], []\n for completion, tid, s in zip(completions, task_ids, seeds):\n gated = CURRICULUM.gate_task(tid)\n action, fmt_ok = parse_action(completion)\n score = run_episode_accumulated(gated, action, seed=int(s))\n rewards.append(score); decisions.append(action.get('decision','?')); fmts.append(fmt_ok)\n CURRICULUM.record(gated, score)\n if METRICS.total_calls % LOG_EVERY == 0:\n print(f' [sample] {gated} s={s} score={score:.3f} {action.get(\"decision\")} fmt={fmt_ok}')\n print(f' curriculum: {CURRICULUM.status()}')\n METRICS.log(rewards, decisions, fmts, list(task_ids))\n if METRICS.step % 5 == 0:\n METRICS.summary()\n return rewards\n\ndef format_reward_fn(completions, **kwargs):\n \"\"\"Format reward: +0.05 valid JSON / -0.05 invalid.\"\"\"\n return [0.15 if parse_action(c)[1] else -0.15 for c in completions]\n\n# Smoke test\nt = env_reward_fn(\n ['{\"decision\":\"APPROVE_FULL\",\"approved_amount\":100.0,\"reason_code\":\"MATCH_CONFIRMED\",\"explanation\":\"Invoice $100.00 matches PO and GRN.\"}'],\n task_id=['easy_perfect_match'], seed=[42]\n)\nprint(f'env_reward smoke: {t[0]:.3f}')\nprint(f'fmt_reward smoke: {format_reward_fn([\"bad json\"])[0]} (expect -0.05)')" ] }, { "cell_type": "code", "execution_count": null, "id": "build_dataset", "metadata": {}, "outputs": [], "source": "from datasets import Dataset\n\n# ALL 10 tasks × 5 seeds = 50 prompts.\n# gate_task() in env_reward_fn redirects locked tasks at score time — not here.\n# This ensures medium/hard task prompts are always in the dataset.\ntask_seed_pairs = [(tid, s) for tid in TRAIN_TASKS for s in range(1, 6)]\nprint(f'[DATASET] {len(task_seed_pairs)} task×seed pairs (all 10 tasks × 5 seeds)')\nprint(f'[CURRICULUM] Currently unlocked: {sorted(CURRICULUM.unlocked)} — gate_task() handles redirection')\n\nrows = []\nfor task_id, seed in task_seed_pairs:\n try:\n reset = requests.post(f'{ENV_URL}/reset', json={'task_id': task_id, 'seed': seed}, timeout=20).json()\n msgs = [{'role':'system','content':SYSTEM_PROMPT},\n {'role':'user','content':obs_to_prompt(reset['observation'])}]\n rows.append({\n 'prompt': tokenizer.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True),\n 'task_id': task_id,\n 'seed': seed,\n })\n except Exception as e:\n print(f' skip {task_id} seed={seed}: {e}')\n\ndataset = Dataset.from_list(rows)\nprint(f'[DATASET] {len(dataset)} prompts ready')" }, { "cell_type": "code", "execution_count": null, "id": "train", "metadata": {}, "outputs": [], "source": [ "from trl import GRPOConfig, GRPOTrainer\nmodel.train()\n\n# per_device_train_batch_size MUST equal num_generations (TRL divisibility requirement)\nconfig = GRPOConfig(\n output_dir = '/tmp/ap_commander_grpo',\n num_train_epochs = NUM_EPOCHS,\n per_device_train_batch_size = NUM_GENERATIONS,\n num_generations = NUM_GENERATIONS,\n gradient_accumulation_steps = 1,\n learning_rate = 2e-5,\n max_completion_length = 250,\n temperature = 0.7,\n logging_steps = 1,\n save_steps = 9999,\n report_to = 'none',\n remove_unused_columns = False,\n beta = 0.1,\n)\ntrainer = GRPOTrainer(\n model = model,\n processing_class = tokenizer,\n reward_funcs = [env_reward_fn, format_reward_fn],\n args = config,\n train_dataset = dataset,\n)\n\nprint(f'Training: {len(dataset)} samples | {NUM_EPOCHS} epochs | {NUM_GENERATIONS} gen/prompt')\nresult = trainer.train()\nprint(f'Done. Loss: {result.training_loss:.4f}')\nMETRICS.summary()" ] }, { "cell_type": "code", "execution_count": null, "id": "post_eval", "metadata": {}, "outputs": [], "source": [ "# Post-training eval — same tasks and seeds as baseline\n", "model.eval()\n", "print('=== POST-TRAINING EVAL ===')\n", "post_results = {}\n", "for task_id in EVAL_TASKS:\n", " scores = []\n", " for seed in EVAL_SEEDS:\n", " score, raw, dec, fmt_ok = eval_one(model, task_id, seed)\n", " scores.append(score)\n", " print(f' {task_id:<35} seed={seed} score={score:.3f} {dec}')\n", " post_results[task_id] = round(sum(scores)/len(scores), 4)\n", "\n", "overall_post = sum(post_results.values()) / len(post_results)\n", "print(f'\\nOverall post-training mean: {overall_post:.3f}')\n", "\n", "print('\\n=== DELTA ===')\n", "for tid in EVAL_TASKS:\n", " b = baseline_results[tid]['mean'] if tid in baseline_results else 0.0\n", " p = post_results[tid]\n", " d = p - b\n", " print(f' {tid:<35} {b:.3f} → {p:.3f} ({d:+.3f})')" ] }, { "cell_type": "code", "execution_count": null, "id": "results_plots", "metadata": {}, "outputs": [], "source": [ "# ── 4-panel results figure (same layout as HF Space) ────────────────────────\n", "fig = dark_fig(figsize=(16, 10))\n", "fig.patch.set_facecolor(BG)\n", "gs = fig.add_gridspec(2, 2, hspace=0.40, wspace=0.30)\n", "\n", "# Panel 1: Before / After eval bars\n", "ax1 = fig.add_subplot(gs[0, 0])\n", "tasks = list(EVAL_TASKS)\n", "short = [t.replace('easy_','').replace('medium_','').replace('hard_','').replace('_',' ').title()\n", " for t in tasks]\n", "xp = np.arange(len(tasks))\n", "b_vals = [baseline_results[t]['mean'] if t in baseline_results else 0.0 for t in tasks]\n", "p_vals = [post_results.get(t, 0.0) for t in tasks]\n", "ax1.bar(xp-0.2, b_vals, 0.35, label='Before GRPO', color='#f85149', alpha=0.85)\n", "ax1.bar(xp+0.2, p_vals, 0.35, label='After GRPO', color='#3fb950', alpha=0.85)\n", "ax1.set_xticks(xp); ax1.set_xticklabels(short, rotation=35, ha='right', fontsize=7)\n", "ax1.set_ylim(0, 1.05); ax1.axhline(0.5, color='#484f58', linestyle='--', alpha=0.6)\n", "ax1.legend(fontsize=8, facecolor=PANEL, edgecolor='#30363d', labelcolor='#c9d1d9')\n", "style_ax(ax1, f'Before vs After — {NUM_EPOCHS} Epochs GRPO')\n", "\n", "# Panel 2: Per-task training mean\n", "ax2 = fig.add_subplot(gs[0, 1])\n", "task_means = {t:round(sum(v)/len(v),3) for t,v in METRICS.reward_by_task.items()}\n", "if task_means:\n", " tm_tasks = list(task_means.keys())\n", " tm_scores = list(task_means.values())\n", " tm_short = [t.replace('easy_','').replace('medium_','').replace('hard_','').replace('_',' ').title()\n", " for t in tm_tasks]\n", " colors = ['#3fb950' if s>=0.7 else '#d29922' if s>=0.4 else '#f85149' for s in tm_scores]\n", " yp2 = range(len(tm_tasks))\n", " ax2.barh(list(yp2), tm_scores, color=colors, alpha=0.85, edgecolor=BG)\n", " ax2.set_yticks(list(yp2)); ax2.set_yticklabels(tm_short, fontsize=7)\n", " ax2.set_xlim(0, 1.05)\n", " ax2.axvline(0.7, color='#3fb950', linestyle='--', linewidth=1, alpha=0.5)\n", " for i,s in enumerate(tm_scores):\n", " ax2.text(s+0.01, i, f'{s:.2f}', va='center', color='#c9d1d9', fontsize=7)\n", "style_ax(ax2, 'Per-Task Training Mean (all seeds)')\n", "\n", "# Panel 3: Decision distribution pie\n", "ax3 = fig.add_subplot(gs[1, 0])\n", "dc = dict(METRICS.decision_counts)\n", "if dc:\n", " colors3 = ['#3fb950','#f85149','#d29922','#a371f7','#58a6ff','#39d353']\n", " wedges, _, autos = ax3.pie(list(dc.values()), labels=None,\n", " autopct='%1.0f%%', colors=colors3[:len(dc)],\n", " startangle=90, pctdistance=0.75,\n", " wedgeprops=dict(edgecolor=BG, linewidth=1.5))\n", " for at in autos: at.set_color(BG); at.set_fontsize(8); at.set_fontweight('bold')\n", " ax3.legend(list(dc.keys()), loc='lower center', bbox_to_anchor=(0.5,-0.15),\n", " ncol=3, fontsize=7, facecolor=PANEL, edgecolor='#30363d', labelcolor='#c9d1d9')\n", "ax3.set_facecolor(PANEL)\n", "ax3.set_title('Decision Distribution', color=TXT, fontsize=10, fontweight='bold', pad=8)\n", "\n", "# Panel 4: Reward curve\n", "ax4 = fig.add_subplot(gs[1, 1])\n", "if METRICS.reward_history:\n", " steps = [s for s,_ in METRICS.reward_history]\n", " rewards = [r for _,r in METRICS.reward_history]\n", " ax4.plot(steps, rewards, color='#58a6ff', alpha=0.30, linewidth=1)\n", " if len(rewards) >= 5:\n", " w = max(3, len(rewards)//15)\n", " sm = np.convolve(rewards, np.ones(w)/w, mode='valid')\n", " ax4.plot(steps[w-1:], sm, color='#79c0ff', linewidth=2, label=f'Smooth w={w}')\n", " mean_r = METRICS.recent_mean()\n", " ax4.axhline(mean_r, color='#f78166', linestyle='--', linewidth=1,\n", " label=f'Recent mean: {mean_r:.3f}')\n", " ax4.set_ylim(0, 1.05)\n", " ax4.legend(fontsize=7, facecolor=PANEL, edgecolor='#30363d', labelcolor='#c9d1d9')\n", "style_ax(ax4, 'Reward Curve', xlabel='Training Step')\n", "\n", "model_short = MODEL_NAME.split('/')[-1]\n", "fig.suptitle(\n", " f'AP Commander GRPO — {model_short} | {NUM_EPOCHS}ep | {NUM_GENERATIONS}gen | '\n", " f'fmt={METRICS.fmt_rate():.1%} | parse_fails={METRICS.parse_failures} | '\n", " f'{datetime.datetime.now().strftime(\"%Y-%m-%d\")}',\n", " color=TXT, fontsize=9, y=0.99\n", ")\n", "plt.savefig('/tmp/results.png', dpi=130, bbox_inches='tight', facecolor=BG)\n", "plt.show()\n", "print('Saved: /tmp/results.png')" ] }, { "cell_type": "code", "execution_count": null, "id": "reward_curve_plot", "metadata": {}, "outputs": [], "source": [ "# Standalone reward curve (for quick progress checks mid-training too)\n", "if METRICS.reward_history:\n", " fig2 = dark_fig(figsize=(12, 4))\n", " ax = fig2.add_subplot(111)\n", " steps = [s for s,_ in METRICS.reward_history]\n", " rewards = [r for _,r in METRICS.reward_history]\n", " ax.plot(steps, rewards, color='#58a6ff', alpha=0.35, linewidth=1, label='Per-step')\n", " if len(rewards) >= 5:\n", " w = max(3, len(rewards)//10)\n", " sm = np.convolve(rewards, np.ones(w)/w, mode='valid')\n", " ax.plot(steps[w-1:], sm, color='#79c0ff', linewidth=2.5, label=f'Smooth w={w}')\n", " ax.axhline(METRICS.recent_mean(), color='#f78166', linestyle='--',\n", " label=f'Recent mean: {METRICS.recent_mean():.3f}')\n", " ax.set_ylim(0, 1.05)\n", " ax.annotate(f'step {steps[-1]} r={rewards[-1]:.3f}',\n", " xy=(steps[-1], rewards[-1]), xytext=(-40,12), textcoords='offset points',\n", " color='#f0f6fc', fontsize=8,\n", " arrowprops=dict(arrowstyle='->', color='#58a6ff', lw=1))\n", " ax.legend(fontsize=8, facecolor=PANEL, edgecolor='#30363d', labelcolor='#c9d1d9')\n", " style_ax(ax, 'Reward Curve', xlabel='Training Step', ylabel='Mean Batch Reward')\n", " fig2.suptitle(f'AP Commander GRPO — {MODEL_NAME.split(\"/\")[-1]}', color=TXT, fontsize=9)\n", " plt.tight_layout()\n", " plt.savefig('/tmp/reward_curve.png', dpi=130, bbox_inches='tight', facecolor=BG)\n", " plt.show()\n", " print('Saved: /tmp/reward_curve.png')" ] }, { "cell_type": "code", "execution_count": null, "id": "save_grpo", "metadata": {}, "outputs": [], "source": [ "import shutil\n", "\n", "model_slug = MODEL_NAME.split('/')[-1].lower().replace('.', '-')\n", "ts = datetime.datetime.now().strftime('%Y-%m-%d_%H%M')\n", "run_dir = f'/tmp/runs/grpo/{model_slug}-{NUM_EPOCHS}ep-{ts}'\n", "os.makedirs(run_dir, exist_ok=True)\n", "\n", "# Save LoRA adapters (guide: save directly, do NOT merge 4-bit naively)\n", "adapter_dir = os.path.join(run_dir, 'adapter')\n", "model.save_pretrained(adapter_dir)\n", "tokenizer.save_pretrained(adapter_dir)\n", "print(f'Adapters saved to {adapter_dir}')\n", "\n", "# Save plots and JSON\n", "for src, dst in [\n", " ('/tmp/results.png', 'results.png'),\n", " ('/tmp/reward_curve.png', 'reward_curve.png'),\n", "]:\n", " if os.path.exists(src):\n", " shutil.copy(src, os.path.join(run_dir, dst))\n", "\n", "output = {\n", " 'timestamp': datetime.datetime.now().isoformat(),\n", " 'run_dir': run_dir,\n", " 'model': MODEL_NAME,\n", " 'epochs': NUM_EPOCHS,\n", " 'num_generations':NUM_GENERATIONS,\n", " 'eval_seeds': EVAL_SEEDS,\n", " 'eval_tasks': EVAL_TASKS,\n", " 'hardware': 'GPU (Colab/HF Spaces)',\n", " 'baseline': {t: baseline_results[t]['mean'] for t in baseline_results} if baseline_results else {},\n", " 'post_training': post_results,\n", " 'delta': {t: round(post_results.get(t,0) - (baseline_results[t]['mean'] if t in baseline_results else 0), 4)\n", " for t in EVAL_TASKS},\n", " 'metrics': {\n", " 'total_reward_calls': METRICS.total_calls,\n", " 'parse_failures': METRICS.parse_failures,\n", " 'format_rate': round(METRICS.fmt_rate(), 4),\n", " 'recent_mean': round(METRICS.recent_mean(), 4),\n", " 'decision_counts': dict(METRICS.decision_counts),\n", " 'per_task_mean': {t: round(sum(v)/len(v),4) for t,v in METRICS.reward_by_task.items()},\n", " },\n", "}\n", "json_path = os.path.join(run_dir, 'training_results.json')\n", "with open(json_path, 'w') as f:\n", " json.dump(output, f, indent=2)\n", "print(f'JSON saved: {json_path}')\n", "\n", "# Upload adapter to HF Hub model repo\n", "if HF_TOKEN:\n", " try:\n", " from huggingface_hub import HfApi\n", " api = HfApi(token=HF_TOKEN)\n", " api.upload_folder(\n", " folder_path=adapter_dir,\n", " repo_id='Pathikreet/ap-commander-adapter',\n", " repo_type='model',\n", " commit_message=f'GRPO {ts} — {model_slug} {NUM_EPOCHS}ep',\n", " )\n", " print('Adapter → HF Hub: Pathikreet/ap-commander-adapter')\n", " except Exception as e:\n", " print(f'Adapter upload skipped: {e}')\n", "\n", " # Upload full run folder to training Space repo\n", " try:\n", " repo_path = f'runs/grpo/{model_slug}-{NUM_EPOCHS}ep-{ts}'\n", " api.upload_folder(\n", " folder_path=run_dir,\n", " path_in_repo=repo_path,\n", " repo_id='Pathikreet/ap-commander-training',\n", " repo_type='space',\n", " commit_message=f'Run artifacts: {os.path.basename(run_dir)}',\n", " ignore_patterns=['adapter/*'],\n", " )\n", " print(f'Run folder → HF Space repo: {repo_path}')\n", " except Exception as e:\n", " print(f'Run folder upload skipped: {e}')\n", "\n", "print(f'\\nAll done. Artifacts in {run_dir}')" ] }, { "cell_type": "markdown", "id": "summary_md", "metadata": {}, "source": [ "## Summary\n", "\n", "| Artifact | Location |\n", "|---|---|\n", "| Baseline results | `runs/baselines/MODEL-DATETIME/baseline_results.json` |\n", "| Baseline plot | `runs/baselines/MODEL-DATETIME/baseline_plot.png` |\n", "| GRPO results | `runs/grpo/MODEL-NEP-DATETIME/training_results.json` |\n", "| GRPO 4-panel plot | `runs/grpo/MODEL-NEP-DATETIME/results.png` |\n", "| Reward curve | `runs/grpo/MODEL-NEP-DATETIME/reward_curve.png` |\n", "| LoRA adapter | `Pathikreet/ap-commander-adapter` on HF Hub |\n", "\n", "All folders are timestamped — re-running never overwrites a previous run." ] } ] }