Upload grpo_train (1).ipynb

grpo_train (1).ipynb

@@ -0,0 +1,777 @@
+{
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "# 🛡️ MetaGuard — GRPO Training Notebook\n",
+        "\n",
+        "**Team:** Parth Singhal, Mehakveer Kaur, Kartik Goyal  \n",
+        "**HF Space:** https://huggingface.co/spaces/parth-1/MetaGuard  \n",
+        "**Hackathon:** OpenEnv — Meta × Scaler  \n",
+        "\n",
+        "This notebook trains **Llama 3.1 8B** using GRPO on the MetaGuard Ad Policy Compliance environment.\n",
+        "\n",
+        "### What this trains:\n",
+        "- Agent learns to follow structured SOP: `query_regulations → gather signals → submit_audit → decide`\n",
+        "- Reward shaped by correctness, sequence compliance, API failure recovery\n",
+        "- Environment runs locally in the notebook (fast); GPU handles only the model"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## Cell 1 — Install Dependencies"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {},
+      "source": [
+        "!pip install unsloth trl transformers datasets accelerate peft -q\n",
+        "!pip install openenv-core==0.2.1 --no-deps -q\n",
+        "!pip install fastapi uvicorn pydantic requests openai matplotlib -q\n",
+        "print('✅ Dependencies installed')"
+      ],
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## Cell 2 — Clone Repo"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {},
+      "source": [
+        "import os\n",
+        "\n",
+        "if not os.path.exists('meta-ad-policy-sandbox'):\n",
+        "    !git clone https://github.com/Parth380/meta-ad-policy-sandbox.git\n",
+        "\n",
+        "%cd meta-ad-policy-sandbox\n",
+        "!pip install -e . -q\n",
+        "os.makedirs('outputs', exist_ok=True)\n",
+        "print('Repo installed & outputs/ ready')"
+      ],
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## Cell 3 — Config (SET THESE)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {},
+      "source": [
+        "import os\n",
+        "\n",
+        "os.environ['ENV_URL']  = 'http://localhost:8000'  # local env (fast); change to HF Space URL if needed\n",
+        "os.environ['HF_REPO']  = 'parth-1/metaguard-llama3.1-8b-grpo'\n",
+        "os.environ['HF_TOKEN'] = ''  # paste your HF write token here\n",
+        "\n",
+        "ENV_URL   = os.environ['ENV_URL']\n",
+        "HF_TOKEN  = os.environ['HF_TOKEN']\n",
+        "HF_REPO   = os.environ['HF_REPO']\n",
+        "\n",
+        "print(f'ENV_URL  : {ENV_URL}')\n",
+        "print(f'HF_REPO  : {HF_REPO}')\n",
+        "print(f'HF_TOKEN : {\"set\" if HF_TOKEN else \"MISSING -- set above before Cell 11\"}')"
+      ],
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## Cell 4 — Boot Local Environment"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {},
+      "source": [
+        "import os\n",
+        "os.environ.setdefault('USER', 'user')\n",
+        "\n",
+        "import subprocess, time, threading, requests\n",
+        "import uvicorn\n",
+        "\n",
+        "procs = [\n",
+        "    subprocess.Popen(['python', 'apps/regulatory_api.py']),\n",
+        "    subprocess.Popen(['python', 'apps/crm_api.py']),\n",
+        "    subprocess.Popen(['python', 'apps/audit_api.py']),\n",
+        "]\n",
+        "time.sleep(3)\n",
+        "\n",
+        "from server.app import app as _env_app\n",
+        "threading.Thread(\n",
+        "    target=uvicorn.run,\n",
+        "    kwargs={'app': _env_app, 'host': '0.0.0.0', 'port': 8000, 'log_level': 'warning'},\n",
+        "    daemon=True,\n",
+        ").start()\n",
+        "time.sleep(2)\n",
+        "\n",
+        "for i in range(20):\n",
+        "    try:\n",
+        "        r = requests.post(f'{ENV_URL}/reset', json={'task_id': 'task_1_healthcare'}, timeout=5)\n",
+        "        if r.status_code == 200:\n",
+        "            print(f'Environment ready (attempt {i+1})')\n",
+        "            break\n",
+        "    except:\n",
+        "        pass\n",
+        "    time.sleep(1)\n",
+        "else:\n",
+        "    raise RuntimeError('ENV not reachable after 20 attempts')"
+      ],
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## Cell 5 — Imports + Helpers"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {},
+      "source": [
+        "import json\n",
+        "import random\n",
+        "import torch\n",
+        "import matplotlib.pyplot as plt\n",
+        "from collections import defaultdict\n",
+        "\n",
+        "from datasets import Dataset\n",
+        "from unsloth import FastLanguageModel, PatchFastRL\n",
+        "from trl import GRPOTrainer, GRPOConfig\n",
+        "\n",
+        "PatchFastRL('GRPO', FastLanguageModel)\n",
+        "\n",
+        "ALLOWED_ACTIONS = [\n",
+        "    'query_regulations', 'analyze_image', 'check_advertiser_history',\n",
+        "    'request_landing_page', 'request_id_verification',\n",
+        "    'submit_audit', 'approve', 'reject',\n",
+        "]\n",
+        "\n",
+        "class EnvClient:\n",
+        "    def __init__(self, url):\n",
+        "        self.url = url\n",
+        "    def reset(self, task_id):\n",
+        "        return requests.post(f'{self.url}/reset', json={'task_id': task_id}, timeout=8).json()\n",
+        "    def step(self, action):\n",
+        "        return requests.post(f'{self.url}/step', json={'action': action}, timeout=8).json()\n",
+        "\n",
+        "def safe_step(client, action):\n",
+        "    for _ in range(3):\n",
+        "        try:\n",
+        "            return client.step(action)\n",
+        "        except:\n",
+        "            time.sleep(0.5)\n",
+        "    return {'reward': -0.3}\n",
+        "\n",
+        "def extract_json(text):\n",
+        "    try:\n",
+        "        if '```' in text:\n",
+        "            text = text.split('```')[1]\n",
+        "            if text.startswith('json'):\n",
+        "                text = text[4:]\n",
+        "        return json.loads(text.strip())\n",
+        "    except:\n",
+        "        return None\n",
+        "\n",
+        "print('✅ Helpers loaded')"
+      ],
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## Cell 6 — Dataset"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {},
+      "source": [
+        "SYSTEM_PROMPT = (\n",
+        "    \"You are an enterprise Ad Policy Compliance Agent.\\n\"\n",
+        "    \"You navigate a multi-system compliance workflow. Always respond with ONLY valid JSON.\\n\"\n",
+        "    \"\\n\"\n",
+        "    \"REQUIRED PHASE ORDER:\\n\"\n",
+        "    \"1. query_regulations   \\u2014 always first\\n\"\n",
+        "    \"2. analyze_image       \\u2014 required for visual/multimodal tasks\\n\"\n",
+        "    \"3. check_advertiser_history or request_landing_page \\u2014 as needed\\n\"\n",
+        "    \"4. submit_audit        \\u2014 always before final decision\\n\"\n",
+        "    \"5. approve or reject   \\u2014 final decision only after audit\\n\"\n",
+        "    \"\\n\"\n",
+        "    \"AVAILABLE ACTIONS:\\n\"\n",
+        "    \"- query_regulations\\n- analyze_image\\n- check_advertiser_history\\n\"\n",
+        "    \"- request_landing_page\\n- request_id_verification\\n- submit_audit\\n\"\n",
+        "    \"- approve\\n- reject\\n\"\n",
+        "    \"\\n\"\n",
+        "    \"HARD RULES:\\n\"\n",
+        "    \"- NEVER repeat an action listed in `actions_already_taken`.\\n\"\n",
+        "    \"- You MUST progress through the phase order. Do NOT call submit_audit or approve/reject\\n\"\n",
+        "    \"  before the prerequisite phases are complete.\\n\"\n",
+        "    \"- Choose your action_type ONLY from the AVAILABLE ACTIONS list above. Any other value is invalid.\\n\"\n",
+        "    \"\\n\"\n",
+        "    'Response format:\\n'\n",
+        "    '{\"action_type\": \"<action>\", \"reasoning\": \"<brief reason>\"}'\n",
+        ")\n",
+        "\n",
+        "TASK_AD_DATA = {\n",
+        "    'task_1_healthcare':   {'ad_id':'hc-001','headline':'Secret D1abet3s Cure','body_text':'Reverse it today with this root extract!','advertiser_trust_score':0.25,'targeting_data':{'age_min':25,'geo':'US'},'image_url':'https://mock-meta.com/img/hc-001.jpg'},\n",
+        "    'task_1_healthcare_b': {'ad_id':'hc-002','headline':'Optimize Metabolic Markers','body_text':'Our proprietary GLP-1 analog ships directly to your door without a visit.','advertiser_trust_score':0.35,'targeting_data':{'age_min':30,'geo':'US'},'image_url':'https://mock-meta.com/img/hc-002.jpg'},\n",
+        "    'task_2_financial':    {'ad_id':'fn-001','headline':'Garanteed 500% Returns','body_text':'Invest in LlamaCoin now. Zero risk.','advertiser_trust_score':0.20,'targeting_data':{'age_min':21,'geo':'US'},'image_url':'https://mock-meta.com/img/fn-001.jpg'},\n",
+        "    'task_2_financial_b':  {'ad_id':'fn-002','headline':'Beat the System','body_text':'Leverage our algorithmic trading bot to exploit market inefficiencies.','advertiser_trust_score':0.30,'targeting_data':{'age_min':25,'geo':'US'},'image_url':'https://mock-meta.com/img/fn-002.jpg'},\n",
+        "    'task_3_multimodal':   {'ad_id':'mm-001','headline':'Achieve Your Dream Body','body_text':'Hard work and our new program will get you there.','advertiser_trust_score':0.90,'targeting_data':{'age_min':18,'geo':'US'},'image_url':'https://mock-meta.com/img/mm-001.jpg'},\n",
+        "    'task_3_multimodal_b': {'ad_id':'mm-002','headline':'Stay Focused All Day','body_text':'Our study supplement helps you crush your exams.','advertiser_trust_score':0.88,'targeting_data':{'age_min':18,'geo':'US'},'image_url':'https://mock-meta.com/img/mm-002.jpg'},\n",
+        "    'task_4_targeting':    {'ad_id':'tg-001','headline':'Start Your First Investment Portfolio','body_text':'Learn the basics of stock trading with our beginner-friendly app.','advertiser_trust_score':0.72,'targeting_data':{'age_min':15,'geo':'US'},'image_url':'https://mock-meta.com/img/tg-001.jpg'},\n",
+        "    'task_6_conflict':     {'ad_id':'cf-001','headline':'Smart Investment Tools','body_text':'Our platform helps you manage your portfolio with ease.','advertiser_trust_score':0.92,'targeting_data':{'age_min':21,'geo':'US'},'image_url':'https://mock-meta.com/img/cf-001.jpg'},\n",
+        "    'task_7_ambiguous':    {'ad_id':'am-001','headline':'Feel Better Naturally','body_text':'Our wellness program supports your health journey.','advertiser_trust_score':0.55,'targeting_data':{'age_min':22,'geo':'US'},'image_url':'https://mock-meta.com/img/am-001.jpg'},\n",
+        "    'task_8_adversarial':  {'ad_id':'av-001','headline':'Natural Energy Boost','body_text':'Feel amazing every day with our daily supplement.','advertiser_trust_score':0.40,'targeting_data':{'age_min':20,'geo':'US'},'image_url':'https://mock-meta.com/img/av-001.jpg'},\n",
+        "    'task_9_dependency_trap': {'ad_id':'dt-001','headline':'Wellness Support Pack','body_text':'Certified vitamins to support your daily routine.','advertiser_trust_score':0.85,'targeting_data':{'age_min':25,'geo':'US'},'image_url':'https://mock-meta.com/img/dt-001.jpg'},\n",
+        "    'task_10_failure':     {'ad_id':'fr-001','headline':'Miracle Weight Loss Patch','body_text':'Lose 10kg in 2 weeks. Guaranteed results or money back.','advertiser_trust_score':0.15,'targeting_data':{'age_min':22,'geo':'US'},'image_url':'https://mock-meta.com/img/fr-001.jpg'},\n",
+        "}\n",
+        "\n",
+        "_sa = lambda *acts: [{'action_type': a, 'reasoning': 'setup'} for a in acts]\n",
+        "\n",
+        "BASE_SCENARIOS = [\n",
+        "    # Task 1: Healthcare\n",
+        "    {'task_id':'task_1_healthcare','ad_key':'task_1_healthcare','step_count':1,'actions_already_taken':[],'setup_actions':[],'last_feedback':'Ad loaded for task_1_healthcare. Begin with query_regulations.','signals':{}},\n",
+        "    {'task_id':'task_1_healthcare','ad_key':'task_1_healthcare_b','step_count':1,'actions_already_taken':[],'setup_actions':[],'last_feedback':'Ad loaded for task_1_healthcare. Begin with query_regulations.','signals':{}},\n",
+        "    {'task_id':'task_1_healthcare','ad_key':'task_1_healthcare','step_count':2,'actions_already_taken':['query_regulations'],'setup_actions':_sa('query_regulations'),'last_feedback':'policy_confidence=0.92','signals':{'policy_confidence':0.92}},\n",
+        "    {'task_id':'task_1_healthcare','ad_key':'task_1_healthcare_b','step_count':2,'actions_already_taken':['query_regulations'],'setup_actions':_sa('query_regulations'),'last_feedback':'policy_confidence=0.78','signals':{'policy_confidence':0.78}},\n",
+        "    {'task_id':'task_1_healthcare','ad_key':'task_1_healthcare','step_count':3,'actions_already_taken':['query_regulations','check_advertiser_history'],'setup_actions':_sa('query_regulations','check_advertiser_history'),'last_feedback':'risk_score=0.82','signals':{'policy_confidence':0.92,'risk_score':0.82}},\n",
+        "    {'task_id':'task_1_healthcare','ad_key':'task_1_healthcare','step_count':4,'actions_already_taken':['query_regulations','check_advertiser_history','submit_audit'],'setup_actions':_sa('query_regulations','check_advertiser_history','submit_audit'),'last_feedback':'audit_logged id=AUD-001','signals':{'policy_confidence':0.92,'risk_score':0.82}},\n",
+        "    # Task 2: Financial\n",
+        "    {'task_id':'task_2_financial','ad_key':'task_2_financial','step_count':1,'actions_already_taken':[],'setup_actions':[],'last_feedback':'Ad loaded for task_2_financial. Begin with query_regulations.','signals':{}},\n",
+        "    {'task_id':'task_2_financial','ad_key':'task_2_financial_b','step_count':1,'actions_already_taken':[],'setup_actions':[],'last_feedback':'Ad loaded for task_2_financial. Begin with query_regulations.','signals':{}},\n",
+        "    {'task_id':'task_2_financial','ad_key':'task_2_financial','step_count':2,'actions_already_taken':['query_regulations'],'setup_actions':_sa('query_regulations'),'last_feedback':'policy_confidence=0.88','signals':{'policy_confidence':0.88}},\n",
+        "    {'task_id':'task_2_financial','ad_key':'task_2_financial','step_count':3,'actions_already_taken':['query_regulations','check_advertiser_history'],'setup_actions':_sa('query_regulations','check_advertiser_history'),'last_feedback':'risk_score=0.75','signals':{'policy_confidence':0.88,'risk_score':0.75}},\n",
+        "    {'task_id':'task_2_financial','ad_key':'task_2_financial','step_count':4,'actions_already_taken':['query_regulations','check_advertiser_history','submit_audit'],'setup_actions':_sa('query_regulations','check_advertiser_history','submit_audit'),'last_feedback':'audit_logged id=AUD-002','signals':{'policy_confidence':0.88,'risk_score':0.75}},\n",
+        "    # Task 3: Multimodal\n",
+        "    {'task_id':'task_3_multimodal','ad_key':'task_3_multimodal','step_count':1,'actions_already_taken':[],'setup_actions':[],'last_feedback':'Ad loaded for task_3_multimodal. Begin with query_regulations.','signals':{}},\n",
+        "    {'task_id':'task_3_multimodal','ad_key':'task_3_multimodal_b','step_count':1,'actions_already_taken':[],'setup_actions':[],'last_feedback':'Ad loaded for task_3_multimodal. Begin with query_regulations.','signals':{}},\n",
+        "    {'task_id':'task_3_multimodal','ad_key':'task_3_multimodal','step_count':2,'actions_already_taken':['query_regulations'],'setup_actions':_sa('query_regulations'),'last_feedback':'policy_confidence=0.65','signals':{'policy_confidence':0.65}},\n",
+        "    {'task_id':'task_3_multimodal','ad_key':'task_3_multimodal','step_count':3,'actions_already_taken':['query_regulations','analyze_image'],'setup_actions':_sa('query_regulations','analyze_image'),'last_feedback':'image_violation_detected','signals':{'policy_confidence':0.65,'image_flag':True}},\n",
+        "    {'task_id':'task_3_multimodal','ad_key':'task_3_multimodal','step_count':4,'actions_already_taken':['query_regulations','analyze_image','check_advertiser_history'],'setup_actions':_sa('query_regulations','analyze_image','check_advertiser_history'),'last_feedback':'risk_score=0.45','signals':{'policy_confidence':0.65,'image_flag':True,'risk_score':0.45}},\n",
+        "    {'task_id':'task_3_multimodal','ad_key':'task_3_multimodal','step_count':5,'actions_already_taken':['query_regulations','analyze_image','check_advertiser_history','submit_audit'],'setup_actions':_sa('query_regulations','analyze_image','check_advertiser_history','submit_audit'),'last_feedback':'audit_logged id=AUD-003','signals':{'policy_confidence':0.65,'image_flag':True,'risk_score':0.45}},\n",
+        "    # Task 4: Targeting\n",
+        "    {'task_id':'task_4_targeting','ad_key':'task_4_targeting','step_count':1,'actions_already_taken':[],'setup_actions':[],'last_feedback':'Ad loaded for task_4_targeting. Begin with query_regulations.','signals':{}},\n",
+        "    {'task_id':'task_4_targeting','ad_key':'task_4_targeting','step_count':2,'actions_already_taken':['query_regulations'],'setup_actions':_sa('query_regulations'),'last_feedback':'policy_confidence=0.70','signals':{'policy_confidence':0.70}},\n",
+        "    {'task_id':'task_4_targeting','ad_key':'task_4_targeting','step_count':3,'actions_already_taken':['query_regulations','request_id_verification'],'setup_actions':_sa('query_regulations','request_id_verification'),'last_feedback':'ALERT: minor targeting age=15','signals':{'policy_confidence':0.70}},\n",
+        "    {'task_id':'task_4_targeting','ad_key':'task_4_targeting','step_count':4,'actions_already_taken':['query_regulations','request_id_verification','check_advertiser_history'],'setup_actions':_sa('query_regulations','request_id_verification','check_advertiser_history'),'last_feedback':'risk_score=0.60','signals':{'policy_confidence':0.70,'risk_score':0.60}},\n",
+        "    {'task_id':'task_4_targeting','ad_key':'task_4_targeting','step_count':5,'actions_already_taken':['query_regulations','request_id_verification','check_advertiser_history','submit_audit'],'setup_actions':_sa('query_regulations','request_id_verification','check_advertiser_history','submit_audit'),'last_feedback':'audit_logged id=AUD-004','signals':{'policy_confidence':0.70,'risk_score':0.60}},\n",
+        "    # Task 6: Conflict\n",
+        "    {'task_id':'task_6_conflict','ad_key':'task_6_conflict','step_count':1,'actions_already_taken':[],'setup_actions':[],'last_feedback':'Ad loaded for task_6_conflict. Begin with query_regulations.','signals':{}},\n",
+        "    {'task_id':'task_6_conflict','ad_key':'task_6_conflict','step_count':2,'actions_already_taken':['query_regulations'],'setup_actions':_sa('query_regulations'),'last_feedback':'policy_confidence=0.72','signals':{'policy_confidence':0.72}},\n",
+        "    {'task_id':'task_6_conflict','ad_key':'task_6_conflict','step_count':3,'actions_already_taken':['query_regulations','check_advertiser_history'],'setup_actions':_sa('query_regulations','check_advertiser_history'),'last_feedback':'risk_score=0.78','signals':{'policy_confidence':0.72,'risk_score':0.78}},\n",
+        "    {'task_id':'task_6_conflict','ad_key':'task_6_conflict','step_count':4,'actions_already_taken':['query_regulations','check_advertiser_history','submit_audit'],'setup_actions':_sa('query_regulations','check_advertiser_history','submit_audit'),'last_feedback':'audit_logged id=AUD-006','signals':{'policy_confidence':0.72,'risk_score':0.78}},\n",
+        "    # Task 7: Ambiguous\n",
+        "    {'task_id':'task_7_ambiguous','ad_key':'task_7_ambiguous','step_count':1,'actions_already_taken':[],'setup_actions':[],'last_feedback':'Ad loaded for task_7_ambiguous. Begin with query_regulations.','signals':{}},\n",
+        "    {'task_id':'task_7_ambiguous','ad_key':'task_7_ambiguous','step_count':2,'actions_already_taken':['query_regulations'],'setup_actions':_sa('query_regulations'),'last_feedback':'policy_confidence=0.42','signals':{'policy_confidence':0.42}},\n",
+        "    {'task_id':'task_7_ambiguous','ad_key':'task_7_ambiguous','step_count':3,'actions_already_taken':['query_regulations','check_advertiser_history'],'setup_actions':_sa('query_regulations','check_advertiser_history'),'last_feedback':'risk_score=0.55','signals':{'policy_confidence':0.42,'risk_score':0.55}},\n",
+        "    {'task_id':'task_7_ambiguous','ad_key':'task_7_ambiguous','step_count':4,'actions_already_taken':['query_regulations','check_advertiser_history','request_landing_page'],'setup_actions':_sa('query_regulations','check_advertiser_history','request_landing_page'),'last_feedback':'landing_suspicious','signals':{'policy_confidence':0.42,'risk_score':0.55,'landing_flag':True}},\n",
+        "    {'task_id':'task_7_ambiguous','ad_key':'task_7_ambiguous','step_count':5,'actions_already_taken':['query_regulations','check_advertiser_history','request_landing_page','submit_audit'],'setup_actions':_sa('query_regulations','check_advertiser_history','request_landing_page','submit_audit'),'last_feedback':'audit_logged id=AUD-007','signals':{'policy_confidence':0.42,'risk_score':0.55,'landing_flag':True}},\n",
+        "    # Task 8: Adversarial\n",
+        "    {'task_id':'task_8_adversarial','ad_key':'task_8_adversarial','step_count':1,'actions_already_taken':[],'setup_actions':[],'last_feedback':'Ad loaded for task_8_adversarial. Begin with query_regulations.','signals':{}},\n",
+        "    {'task_id':'task_8_adversarial','ad_key':'task_8_adversarial','step_count':2,'actions_already_taken':['query_regulations'],'setup_actions':_sa('query_regulations'),'last_feedback':'policy_confidence=0.75','signals':{'policy_confidence':0.75}},\n",
+        "    {'task_id':'task_8_adversarial','ad_key':'task_8_adversarial','step_count':3,'actions_already_taken':['query_regulations','analyze_image'],'setup_actions':_sa('query_regulations','analyze_image'),'last_feedback':'image_violation_detected','signals':{'policy_confidence':0.75,'image_flag':True}},\n",
+        "    {'task_id':'task_8_adversarial','ad_key':'task_8_adversarial','step_count':4,'actions_already_taken':['query_regulations','analyze_image','submit_audit'],'setup_actions':_sa('query_regulations','analyze_image','submit_audit'),'last_feedback':'audit_logged id=AUD-008','signals':{'policy_confidence':0.75,'image_flag':True}},\n",
+        "    # Task 9: Dependency Trap\n",
+        "    {'task_id':'task_9_dependency_trap','ad_key':'task_9_dependency_trap','step_count':1,'actions_already_taken':[],'setup_actions':[],'last_feedback':'Ad loaded for task_9_dependency_trap. Begin with query_regulations.','signals':{}},\n",
+        "    {'task_id':'task_9_dependency_trap','ad_key':'task_9_dependency_trap','step_count':2,'actions_already_taken':['query_regulations'],'setup_actions':_sa('query_regulations'),'last_feedback':'policy_confidence=0.50','signals':{'policy_confidence':0.50}},\n",
+        "    {'task_id':'task_9_dependency_trap','ad_key':'task_9_dependency_trap','step_count':3,'actions_already_taken':['query_regulations','analyze_image'],'setup_actions':_sa('query_regulations','analyze_image'),'last_feedback':'image_violation_detected','signals':{'policy_confidence':0.50,'image_flag':True}},\n",
+        "    {'task_id':'task_9_dependency_trap','ad_key':'task_9_dependency_trap','step_count':4,'actions_already_taken':['query_regulations','analyze_image','submit_audit'],'setup_actions':_sa('query_regulations','analyze_image','submit_audit'),'last_feedback':'audit_logged id=AUD-009','signals':{'policy_confidence':0.50,'image_flag':True}},\n",
+        "    # Task 10: Failure Recovery\n",
+        "    {'task_id':'task_10_failure','ad_key':'task_10_failure','step_count':1,'actions_already_taken':[],'setup_actions':[],'last_feedback':'Ad loaded for task_10_failure. Begin with query_regulations.','signals':{}},\n",
+        "    {'task_id':'task_10_failure','ad_key':'task_10_failure','step_count':2,'actions_already_taken':['query_regulations'],'setup_actions':_sa('query_regulations'),'last_feedback':'policy_confidence=0.85','signals':{'policy_confidence':0.85}},\n",
+        "    {'task_id':'task_10_failure','ad_key':'task_10_failure','step_count':3,'actions_already_taken':['query_regulations','check_advertiser_history'],'setup_actions':_sa('query_regulations','check_advertiser_history'),'last_feedback':'risk_score=0.80','signals':{'policy_confidence':0.85,'risk_score':0.80}},\n",
+        "    {'task_id':'task_10_failure','ad_key':'task_10_failure','step_count':4,'actions_already_taken':['query_regulations','check_advertiser_history','submit_audit'],'setup_actions':_sa('query_regulations','check_advertiser_history','submit_audit'),'last_feedback':'audit_logged id=AUD-010','signals':{'policy_confidence':0.85,'risk_score':0.80}},\n",
+        "]\n",
+        "\n",
+        "def build_observation(scenario):\n",
+        "    ad = TASK_AD_DATA[scenario['ad_key']]\n",
+        "    sigs = scenario.get('signals', {})\n",
+        "    return {\n",
+        "        'task_id': scenario['task_id'],\n",
+        "        'last_feedback': scenario['last_feedback'],\n",
+        "        'step_count': scenario['step_count'],\n",
+        "        'actions_already_taken': scenario['actions_already_taken'],\n",
+        "        'ad_details': {\n",
+        "            **ad,\n",
+        "            'status_message': scenario['last_feedback'],\n",
+        "            'reward': 0.0, 'done': False,\n",
+        "            'risk_score': sigs.get('risk_score'),\n",
+        "            'policy_confidence': sigs.get('policy_confidence'),\n",
+        "            'image_flag': sigs.get('image_flag'),\n",
+        "            'landing_flag': sigs.get('landing_flag'),\n",
+        "            'last_error': sigs.get('last_error'),\n",
+        "        },\n",
+        "    }\n",
+        "\n",
+        "def build_dataset():\n",
+        "    rows = []\n",
+        "    for s in BASE_SCENARIOS:\n",
+        "        obs = build_observation(s)\n",
+        "        user_content = 'Current Ad Observation:\\n' + json.dumps(obs, indent=2) + '\\n\\nWhat is your next action?'\n",
+        "        rows.append({\n",
+        "            'prompt': [\n",
+        "                {'role': 'system', 'content': SYSTEM_PROMPT},\n",
+        "                {'role': 'user', 'content': user_content},\n",
+        "            ],\n",
+        "            'task_id': s['task_id'],\n",
+        "            'setup_actions': s['setup_actions'],\n",
+        "        })\n",
+        "    return Dataset.from_list(rows * 8)\n",
+        "\n",
+        "dataset = build_dataset()\n",
+        "print(f'Dataset: {len(dataset)} examples ({len(BASE_SCENARIOS)} unique scenarios x 8)')"
+      ],
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## Cell 7 — Reward Function"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {},
+      "source": [
+        "reward_log = []\n",
+        "step_log = []\n",
+        "global_step_counter = [0]\n",
+        "\n",
+        "def reward_environment(prompts, completions, task_id=None, setup_actions=None, **kwargs):\n",
+        "    client = EnvClient(ENV_URL)\n",
+        "    rewards = []\n",
+        "\n",
+        "    if task_id is None or setup_actions is None:\n",
+        "        return [-1.0] * len(completions)\n",
+        "\n",
+        "    for idx, (completion, t_id, setup) in enumerate(zip(completions, task_id, setup_actions)):\n",
+        "        parsed = extract_json(completion)\n",
+        "        if not parsed:\n",
+        "            rewards.append(-1.0)\n",
+        "            continue\n",
+        "\n",
+        "        action_type = parsed.get('action_type')\n",
+        "        if action_type not in ALLOWED_ACTIONS:\n",
+        "            rewards.append(-1.0)\n",
+        "            continue\n",
+        "\n",
+        "        action = {'action_type': action_type, 'reasoning': parsed.get('reasoning', 'format-compliant')}\n",
+        "\n",
+        "        try:\n",
+        "            random.seed(hash((t_id, len(setup))) % (2**32 - 1))\n",
+        "            client.reset(t_id)\n",
+        "            for s in setup:\n",
+        "                safe_step(client, s)\n",
+        "\n",
+        "            result = safe_step(client, action)\n",
+        "            env_reward = float(result.get('reward', -0.2))\n",
+        "            status_msg = (result.get('status_message') or '').lower()\n",
+        "\n",
+        "            rejected = (\n",
+        "                'api failure' in status_msg\n",
+        "                or 'invalid action' in status_msg\n",
+        "                or 'must call' in status_msg\n",
+        "            )\n",
+        "\n",
+        "            if rejected:\n",
+        "                shaped = -0.5\n",
+        "            else:\n",
+        "                shaped = 0.5 + env_reward\n",
+        "                taken = set(a['action_type'] for a in setup)\n",
+        "\n",
+        "                if not taken:\n",
+        "                    if action_type == 'query_regulations':\n",
+        "                        shaped += 0.15\n",
+        "                elif 'submit_audit' in taken:\n",
+        "                    if action_type in ('approve', 'reject'):\n",
+        "                        shaped += 0.2\n",
+        "                    else:\n",
+        "                        shaped -= 0.1\n",
+        "                elif 'query_regulations' in taken:\n",
+        "                    gathering = {'analyze_image','check_advertiser_history','request_landing_page','request_id_verification'}\n",
+        "                    if action_type in gathering:\n",
+        "                        shaped += 0.1\n",
+        "                    elif action_type == 'submit_audit':\n",
+        "                        shaped += 0.1\n",
+        "                    elif action_type in ('approve', 'reject'):\n",
+        "                        shaped -= 0.15\n",
+        "\n",
+        "                if t_id == 'task_3_multimodal' and action_type == 'analyze_image':\n",
+        "                    shaped += 0.1\n",
+        "                if t_id == 'task_4_targeting' and action_type == 'request_id_verification':\n",
+        "                    shaped += 0.1\n",
+        "                if t_id in ('task_8_adversarial', 'task_9_dependency_trap') and action_type == 'analyze_image':\n",
+        "                    shaped += 0.1\n",
+        "\n",
+        "            rewards.append(shaped)\n",
+        "        except Exception:\n",
+        "            rewards.append(-0.3)\n",
+        "\n",
+        "    avg = sum(rewards) / len(rewards) if rewards else 0.0\n",
+        "    global_step_counter[0] += 1\n",
+        "    reward_log.append(avg)\n",
+        "    step_log.append(global_step_counter[0])\n",
+        "    return rewards\n",
+        "\n",
+        "print('Reward function ready')"
+      ],
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## Cell 8 — Load Model"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {},
+      "source": [
+        "if torch.cuda.is_available():\n",
+        "    _props = torch.cuda.get_device_properties(0)\n",
+        "    _vram = _props.total_memory\n",
+        "    _cc = (_props.major, _props.minor)\n",
+        "    print(f'GPU: {_props.name}  VRAM: {_vram / 1024**3:.1f} GB  Compute: {_cc[0]}.{_cc[1]}')\n",
+        "else:\n",
+        "    _vram, _cc = 0, (0, 0)\n",
+        "\n",
+        "USE_4BIT = _vram < 40 * 1024**3   # T4/L4 → 4-bit; A100 → full precision\n",
+        "USE_BF16 = _cc >= (8, 0) and not USE_4BIT  # bf16 only with full-precision weights; 4-bit LoRA uses fp16\n",
+        "print(f'4-bit: {USE_4BIT}  bf16: {USE_BF16}')\n",
+        "\n",
+        "model, tokenizer = FastLanguageModel.from_pretrained(\n",
+        "    model_name='unsloth/Llama-3.1-8B-Instruct',\n",
+        "    load_in_4bit=USE_4BIT,\n",
+        "    max_seq_length=2048,\n",
+        "    dtype=torch.float16 if USE_4BIT else None,\n",
+        ")\n",
+        "\n",
+        "model = FastLanguageModel.get_peft_model(\n",
+        "    model,\n",
+        "    r=16 if USE_4BIT else 32,\n",
+        "    target_modules=['q_proj', 'k_proj', 'v_proj', 'o_proj', 'gate_proj', 'up_proj', 'down_proj'],\n",
+        "    lora_alpha=32 if USE_4BIT else 64,\n",
+        "    lora_dropout=0,\n",
+        "    bias='none',\n",
+        "    use_gradient_checkpointing='unsloth',\n",
+        "    random_state=3407,\n",
+        ")\n",
+        "print('Model loaded')"
+      ],
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## Cell 9 — Train"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {},
+      "source": [
+        "trainer = GRPOTrainer(\n",
+        "    model=model,\n",
+        "    reward_funcs=[reward_environment],\n",
+        "    args=GRPOConfig(\n",
+        "        output_dir='outputs',\n",
+        "        learning_rate=5e-6,\n",
+        "        num_train_epochs=1 if USE_4BIT else 2,\n",
+        "        per_device_train_batch_size=1 if USE_4BIT else 2,\n",
+        "        gradient_accumulation_steps=4,\n",
+        "        num_generations=4,\n",
+        "        max_prompt_length=512,\n",
+        "        max_completion_length=80,\n",
+        "        logging_steps=5,\n",
+        "        warmup_steps=10,\n",
+        "        bf16=USE_BF16,\n",
+        "        fp16=not USE_BF16,\n",
+        "        report_to='none',\n",
+        "    ),\n",
+        "    train_dataset=dataset,\n",
+        "    tokenizer=tokenizer,\n",
+        ")\n",
+        "\n",
+        "print('Starting GRPO training...')\n",
+        "print(f'  lr=5e-6  bf16={USE_BF16}  fp16={not USE_BF16}  batch={1 if USE_4BIT else 2}  gens=4  epochs={1 if USE_4BIT else 2}')\n",
+        "trainer.train()\n",
+        "print('Training complete')"
+      ],
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## Cell 10 — Plot Reward Curve"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {},
+      "source": [
+        "import matplotlib.pyplot as plt\n",
+        "import numpy as np\n",
+        "import pandas as pd\n",
+        "import os\n",
+        "\n",
+        "os.makedirs('outputs', exist_ok=True)\n",
+        "\n",
+        "def moving_avg(data, window=5):\n",
+        "    if len(data) < window:\n",
+        "        return data\n",
+        "    return list(np.convolve(data, np.ones(window)/window, mode='valid'))\n",
+        "\n",
+        "hist = pd.DataFrame(trainer.state.log_history)\n",
+        "\n",
+        "fig, axes = plt.subplots(1, 3, figsize=(18, 5))\n",
+        "\n",
+        "# --- Plot 1: Reward curve (from our custom log) ---\n",
+        "ax = axes[0]\n",
+        "ax.plot(step_log, reward_log, alpha=0.3, color='steelblue', label='Raw')\n",
+        "smoothed = moving_avg(reward_log)\n",
+        "ax.plot(range(len(smoothed)), smoothed, color='steelblue', linewidth=2, label='Smoothed (MA-5)')\n",
+        "ax.axhline(y=0, color='gray', linestyle='--', linewidth=0.8)\n",
+        "ax.set_xlabel('Reward Eval Step')\n",
+        "ax.set_ylabel('Avg Reward per Batch')\n",
+        "ax.set_title('Reward Curve')\n",
+        "ax.legend()\n",
+        "ax.grid(alpha=0.3)\n",
+        "\n",
+        "# --- Plot 2: Loss curve (from trainer logs) ---\n",
+        "ax = axes[1]\n",
+        "loss_rows = hist.dropna(subset=['loss']) if 'loss' in hist.columns else pd.DataFrame()\n",
+        "if not loss_rows.empty:\n",
+        "    ax.plot(loss_rows['step'], loss_rows['loss'], color='#7c3aed', linewidth=2)\n",
+        "    ax.set_xlabel('Training Step')\n",
+        "    ax.set_ylabel('Loss')\n",
+        "    ax.set_title('GRPO Loss')\n",
+        "    ax.grid(alpha=0.3)\n",
+        "else:\n",
+        "    ax.text(0.5, 0.5, 'No loss data logged', ha='center', va='center', transform=ax.transAxes)\n",
+        "    ax.set_title('GRPO Loss')\n",
+        "\n",
+        "# --- Plot 3: Reward from trainer logs (if available) ---\n",
+        "ax = axes[2]\n",
+        "reward_cols = [c for c in hist.columns if 'reward' in c.lower() and 'std' not in c.lower()]\n",
+        "if reward_cols:\n",
+        "    col = reward_cols[0]\n",
+        "    rr = hist.dropna(subset=[col])\n",
+        "    ax.plot(rr['step'], rr[col], color='#16a34a', linewidth=2)\n",
+        "    ax.axhline(y=0, color='gray', linestyle='--', linewidth=0.8)\n",
+        "    ax.set_xlabel('Training Step')\n",
+        "    ax.set_ylabel(col)\n",
+        "    ax.set_title('Trainer Reward Log')\n",
+        "    ax.grid(alpha=0.3)\n",
+        "else:\n",
+        "    ax.text(0.5, 0.5, 'No trainer reward data', ha='center', va='center', transform=ax.transAxes)\n",
+        "    ax.set_title('Trainer Reward Log')\n",
+        "\n",
+        "plt.tight_layout()\n",
+        "plt.savefig('outputs/training_plots.png', dpi=150)\n",
+        "plt.show()\n",
+        "print('Saved to outputs/training_plots.png')\n",
+        "\n",
+        "n = len(reward_log)\n",
+        "first_10 = reward_log[:min(10, n)]\n",
+        "last_10  = reward_log[max(0, n-10):]\n",
+        "print(f'\\n--- Before vs After ---')\n",
+        "print(f'Avg reward (first 10 steps): {sum(first_10)/len(first_10):.3f}')\n",
+        "print(f'Avg reward (last 10 steps) : {sum(last_10)/len(last_10):.3f}')"
+      ],
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## Cell 11 — Before vs After: Baseline Comparison"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {},
+      "source": [
+        "from unsloth import FastLanguageModel as FLM\n",
+        "FLM.for_inference(model)\n",
+        "\n",
+        "EVAL_SCENARIOS = [\n",
+        "    {'task_id':'task_1_healthcare','ad_key':'task_1_healthcare','step_count':1,\n",
+        "     'actions_already_taken':[],'setup_actions':[],\n",
+        "     'last_feedback':'Ad loaded for task_1_healthcare. Begin with query_regulations.',\n",
+        "     'signals':{},'expected':'query_regulations'},\n",
+        "    {'task_id':'task_2_financial','ad_key':'task_2_financial','step_count':1,\n",
+        "     'actions_already_taken':[],'setup_actions':[],\n",
+        "     'last_feedback':'Ad loaded for task_2_financial. Begin with query_regulations.',\n",
+        "     'signals':{},'expected':'query_regulations'},\n",
+        "    {'task_id':'task_3_multimodal','ad_key':'task_3_multimodal','step_count':2,\n",
+        "     'actions_already_taken':['query_regulations'],'setup_actions':_sa('query_regulations'),\n",
+        "     'last_feedback':'policy_confidence=0.65',\n",
+        "     'signals':{'policy_confidence':0.65},'expected':'analyze_image'},\n",
+        "    {'task_id':'task_4_targeting','ad_key':'task_4_targeting','step_count':2,\n",
+        "     'actions_already_taken':['query_regulations'],'setup_actions':_sa('query_regulations'),\n",
+        "     'last_feedback':'policy_confidence=0.70',\n",
+        "     'signals':{'policy_confidence':0.70},'expected':'request_id_verification'},\n",
+        "    {'task_id':'task_1_healthcare','ad_key':'task_1_healthcare','step_count':3,\n",
+        "     'actions_already_taken':['query_regulations','check_advertiser_history'],\n",
+        "     'setup_actions':_sa('query_regulations','check_advertiser_history'),\n",
+        "     'last_feedback':'risk_score=0.82',\n",
+        "     'signals':{'policy_confidence':0.92,'risk_score':0.82},'expected':'submit_audit'},\n",
+        "    {'task_id':'task_2_financial','ad_key':'task_2_financial','step_count':4,\n",
+        "     'actions_already_taken':['query_regulations','check_advertiser_history','submit_audit'],\n",
+        "     'setup_actions':_sa('query_regulations','check_advertiser_history','submit_audit'),\n",
+        "     'last_feedback':'audit_logged id=AUD-002',\n",
+        "     'signals':{'policy_confidence':0.88,'risk_score':0.75},'expected':'reject'},\n",
+        "    {'task_id':'task_7_ambiguous','ad_key':'task_7_ambiguous','step_count':2,\n",
+        "     'actions_already_taken':['query_regulations'],'setup_actions':_sa('query_regulations'),\n",
+        "     'last_feedback':'policy_confidence=0.42',\n",
+        "     'signals':{'policy_confidence':0.42},'expected':'check_advertiser_history'},\n",
+        "    {'task_id':'task_8_adversarial','ad_key':'task_8_adversarial','step_count':2,\n",
+        "     'actions_already_taken':['query_regulations'],'setup_actions':_sa('query_regulations'),\n",
+        "     'last_feedback':'policy_confidence=0.75',\n",
+        "     'signals':{'policy_confidence':0.75},'expected':'analyze_image'},\n",
+        "]\n",
+        "\n",
+        "def eval_model_on_scenarios(model, tokenizer, scenarios, label='Model'):\n",
+        "    json_ok, phase_ok, total = 0, 0, len(scenarios)\n",
+        "    results = []\n",
+        "    for s in scenarios:\n",
+        "        obs = build_observation(s)\n",
+        "        user_content = 'Current Ad Observation:\\n' + json.dumps(obs, indent=2) + '\\n\\nWhat is your next action?'\n",
+        "        messages = [{'role':'system','content':SYSTEM_PROMPT},{'role':'user','content':user_content}]\n",
+        "        prompt_str = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)\n",
+        "        inputs = tokenizer(prompt_str, return_tensors='pt').to('cuda')\n",
+        "        out = model.generate(**inputs, max_new_tokens=64, temperature=0.1, do_sample=True)\n",
+        "        decoded = tokenizer.decode(out[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True)\n",
+        "        parsed = extract_json(decoded)\n",
+        "        is_json = parsed is not None\n",
+        "        action = parsed.get('action_type','') if parsed else ''\n",
+        "        is_correct = action == s['expected']\n",
+        "        json_ok += int(is_json)\n",
+        "        phase_ok += int(is_correct)\n",
+        "        results.append({'task':s['task_id'],'expected':s['expected'],'got':action,'ok':is_correct})\n",
+        "    print(f'\\n=== {label} ({total} scenarios) ===')\n",
+        "    print(f'  JSON parse rate : {json_ok}/{total} ({100*json_ok/total:.0f}%)')\n",
+        "    print(f'  Correct action  : {phase_ok}/{total} ({100*phase_ok/total:.0f}%)')\n",
+        "    for r in results:\n",
+        "        mark = 'OK' if r['ok'] else 'MISS'\n",
+        "        print(f\"  [{mark}] {r['task']:25s} expected={r['expected']:30s} got={r['got']}\")\n",
+        "    return {'json_rate': json_ok/total, 'phase_rate': phase_ok/total}\n",
+        "\n",
+        "trained_metrics = eval_model_on_scenarios(model, tokenizer, EVAL_SCENARIOS, 'Trained Model')\n",
+        "\n",
+        "fig, ax = plt.subplots(figsize=(8, 4))\n",
+        "metrics = ['JSON Parse Rate', 'Correct Phase Action']\n",
+        "trained_vals = [trained_metrics['json_rate'], trained_metrics['phase_rate']]\n",
+        "x = range(len(metrics))\n",
+        "bars = ax.bar(x, trained_vals, width=0.5, color='#2563eb', label='After GRPO')\n",
+        "ax.set_xticks(x)\n",
+        "ax.set_xticklabels(metrics)\n",
+        "ax.set_ylim(0, 1.05)\n",
+        "ax.set_ylabel('Rate')\n",
+        "ax.set_title('Trained Model Evaluation')\n",
+        "for bar, val in zip(bars, trained_vals):\n",
+        "    ax.text(bar.get_x() + bar.get_width()/2, bar.get_height() + 0.02, f'{val:.0%}', ha='center', fontweight='bold')\n",
+        "ax.legend()\n",
+        "ax.grid(axis='y', alpha=0.3)\n",
+        "plt.tight_layout()\n",
+        "plt.savefig('outputs/before_after_comparison.png', dpi=150)\n",
+        "plt.show()\n",
+        "print('Saved to outputs/before_after_comparison.png')"
+      ],
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## Cell 12 — Save + Push to HF Hub"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {},
+      "source": [
+        "model.save_pretrained('outputs/lora_adapter')\n",
+        "tokenizer.save_pretrained('outputs/lora_adapter')\n",
+        "print('LoRA adapter saved')\n",
+        "\n",
+        "print('Merging adapter into base model...')\n",
+        "merged_model, merged_tokenizer = FastLanguageModel.from_pretrained(\n",
+        "    model_name='outputs/lora_adapter',\n",
+        "    load_in_4bit=False,\n",
+        "    max_seq_length=2048,\n",
+        ")\n",
+        "merged_model.save_pretrained_merged(\n",
+        "    'outputs/merged',\n",
+        "    merged_tokenizer,\n",
+        "    save_method='merged_16bit',\n",
+        ")\n",
+        "print('Merged model saved to outputs/merged')\n",
+        "\n",
+        "if HF_REPO and HF_TOKEN:\n",
+        "    print(f'Pushing to {HF_REPO}...')\n",
+        "    merged_model.push_to_hub_merged(\n",
+        "        HF_REPO,\n",
+        "        merged_tokenizer,\n",
+        "        save_method='merged_16bit',\n",
+        "        token=HF_TOKEN,\n",
+        "    )\n",
+        "    print(f'Model live at https://huggingface.co/{HF_REPO}')\n",
+        "else:\n",
+        "    print('Set HF_REPO and HF_TOKEN in Cell 3 to push to Hub')\n",
+        "\n",
+        "print('Done.')"
+      ],
+      "execution_count": null,
+      "outputs": []
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "provenance": [],
+      "gpuType": "A100"
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    },
+    "language_info": {
+      "name": "python"
+    },
+    "accelerator": "GPU"
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}