Add V4.2 GRPO training notebook (Gold Standard, 0.5B)
Browse files
notebooks/v4_2_instruct_grpo.ipynb
ADDED
|
@@ -0,0 +1,214 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"cells": [
|
| 3 |
+
{
|
| 4 |
+
"cell_type": "markdown",
|
| 5 |
+
"metadata": {},
|
| 6 |
+
"source": "# Tucano2 Commerce β GRPO Training V4.2 (Gold Standard, 0.5B)\n\n**Reference:** `docs/v4_2-handoff.md` \n**Base:** V4.1 notebook with 8 targeted changes\n\n## V4.2 Changes from V4.1\n\n| # | Change | V4.1 | V4.2 | Why |\n|---|--------|------|------|-----|\n| 1 | Eval suite | 15 mixed samples | **65 stratified** (20 ext + 15 sql + 15 ins + 15 push) | Insights swing Β±0.22 was eval noise on nβ2 |\n| 2 | Reward audit | None | **Spearman Ο > 0.70 gate** (20 completions, human-scored) | Parser bug persisted 3 versions |\n| 3 | SQL reward | Heuristic vocabulary | **Validation-aware** (SQL syntax + query/explanation + numerics + domain) | SQL stagnant +3.8% β reward was ceiling |\n| 4 | Max steps | 600 | **1,500** (~2.5 epochs) | Only 40% data seen; eval still improving at step 500 |\n| 5 | GDPO normalization | Batch-level reward | **Per-component normalization** before aggregation | GDPO Β§3.1: preserves 4Γ more advantage groups |\n| 6 | Task weighting | Equal (0.25 each) | **Dynamic IWU** (upweight stagnating tasks) | MT-GRPO Β§3.2: prevents easy-task collapse |\n| 7 | Seeds | Single run (42) | **3 seeds** (42, 123, 456) with reported CIs | Minimum for credible ML result |\n| 8 | Best checkpoint | Save at end | **Explicit best_checkpoint/** on eval improvement | GRPOTrainer lacks load_best_model_at_end |\n\n**Prerequisites:**\n- Upload `data/pairs/train.jsonl` and `data/pairs/eval.jsonl` to `./data/pairs/`\n- Hardware: L4 (24GB), PyTorch kernel, bf16 supported\n- Estimated runtime: ~12h per seed (1,500 steps Γ ~30s/step)\n- Run 3 times, changing only `CURRENT_SEED` in Cell 3\n\n---\n\n*V4.2 is the last 0.5B run. Its purpose is not to find more improvement β it is to know exactly what was found and why, with enough statistical rigor to say so in writing.*"
|
| 7 |
+
},
|
| 8 |
+
{
|
| 9 |
+
"cell_type": "markdown",
|
| 10 |
+
"metadata": {},
|
| 11 |
+
"source": "---\n\n## Cell 1: Dependencies\n\n**Gate:** No errors. Verify TRL 0.24.0 installed.\n\n**V4.2 additions:** `scipy` for Spearman Ο in reward audit."
|
| 12 |
+
},
|
| 13 |
+
{
|
| 14 |
+
"cell_type": "code",
|
| 15 |
+
"execution_count": null,
|
| 16 |
+
"metadata": {},
|
| 17 |
+
"outputs": [],
|
| 18 |
+
"source": "# Cell 1 β Clean install\n# Run after kernel restart\n\n!pip install \"unsloth\"\n!pip install \"trl==0.24.0\" --no-deps\n!pip install \"rich\" \"wandb\"\n!pip install \"json-repair\" # V4.1: robust JSON parser for Portuguese LLM output\n!pip install \"scipy\" # V4.2: Spearman Ο for reward audit"
|
| 19 |
+
},
|
| 20 |
+
{
|
| 21 |
+
"cell_type": "markdown",
|
| 22 |
+
"metadata": {},
|
| 23 |
+
"source": "---\n\n## Cell 2: GPU + Unsloth Verification\n\n**Gate:** CUDA available, bf16=True, VRAM > 20GB, TRL 0.24.0."
|
| 24 |
+
},
|
| 25 |
+
{
|
| 26 |
+
"cell_type": "code",
|
| 27 |
+
"execution_count": null,
|
| 28 |
+
"metadata": {},
|
| 29 |
+
"outputs": [],
|
| 30 |
+
"source": "import torch\n\nprint(f\"CUDA available: {torch.cuda.is_available()}\")\nprint(f\"GPU: {torch.cuda.get_device_name(0)}\")\nprint(f\"VRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB\")\nprint(f\"bf16 support: {torch.cuda.is_bf16_supported()}\")\n\nfrom unsloth import FastLanguageModel\nprint(f\"\\nβ Unsloth loaded\")\n\nimport trl\nassert trl.__version__ == \"0.24.0\", f\"Expected TRL 0.24.0, got {trl.__version__}\"\nprint(f\"β TRL {trl.__version__}\")\n\nimport transformers\nprint(f\"β Transformers {transformers.__version__}\")"
|
| 31 |
+
},
|
| 32 |
+
{
|
| 33 |
+
"cell_type": "markdown",
|
| 34 |
+
"metadata": {},
|
| 35 |
+
"source": "---\n\n## Cell 3: Config Constants\n\n**V4.2 changes:**\n- `MAX_STEPS`: 600 β **1,500** (multi-epoch, ~2.5Γ full dataset)\n- `EVAL_STEPS`: 20 β **50** (more frequent eval relative to epoch boundaries)\n- `SAVE_STEPS`: 50 β **100** (scaled for longer run)\n- `SEEDS`: Added multi-seed support. Change `CURRENT_SEED` per run.\n- `EVAL_TOTAL = 65`: Stratified eval set (20 ext + 15 sql + 15 ins + 15 push)\n\n**Everything else UNCHANGED from V4.1** (validated config)."
|
| 36 |
+
},
|
| 37 |
+
{
|
| 38 |
+
"cell_type": "code",
|
| 39 |
+
"execution_count": null,
|
| 40 |
+
"metadata": {},
|
| 41 |
+
"outputs": [],
|
| 42 |
+
"source": "import os\nimport json\nimport re\nimport time\nimport random\nimport gc\nimport math\nimport warnings\nfrom pathlib import Path\n\n# ββ Suppress noisy deprecation warnings from Transformers 5.5.0 ββββββββββββββ\nwarnings.filterwarnings(\"ignore\", message=\".*AttentionMaskConverter.*\")\nwarnings.filterwarnings(\"ignore\", message=\".*Passing `generation_config` together with.*\")\nwarnings.filterwarnings(\"ignore\", message=\".*max_new_tokens.*max_length.*\")\nwarnings.filterwarnings(\"ignore\", category=FutureWarning)\n\n# ββ Disable Unsloth kernel recompilation βββββββββββββββββββββββββββββββββββββ\nos.environ[\"UNSLOTH_COMPILE_DISABLE\"] = \"1\"\nos.environ[\"PYTORCH_CUDA_ALLOC_CONF\"] = \"expandable_segments:True\"\n\n# ββ V4.2: Multi-seed support ββββββββββββββββββββββββββββββββββββββββββββββββ\nSEEDS = [42, 123, 456]\nCURRENT_SEED = 42 # β CHANGE THIS PER RUN (42, 123, 456)\n\n# Set all random seeds\nrandom.seed(CURRENT_SEED)\ntorch.manual_seed(CURRENT_SEED)\nif torch.cuda.is_available():\n torch.cuda.manual_seed_all(CURRENT_SEED)\n\n# ββ Model ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ\nMODEL_ID = \"Polygl0t/Tucano2-qwen-0.5B-Instruct\"\nMAX_SEQ_LENGTH = 2048 # model supports 4096, but 2048 is plenty for Instruct (no <think> overhead)\nMODELS_DIR = Path(\"/home/jupyter/tucano2/models\")\nADAPTER_DIR = MODELS_DIR / f\"tucano2-0.5B-instruct-grpo-v4.2-seed{CURRENT_SEED}\"\nCHECKPOINT_DIR = ADAPTER_DIR / \"checkpoints\"\n\n# ββ Data βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ\nDATA_DIR = Path(\"/home/jupyter/tucano2/data\")\nTRAIN_FILE = DATA_DIR / \"pairs\" / \"train.jsonl\"\nEVAL_FILE = DATA_DIR / \"pairs\" / \"eval.jsonl\" # V4.2: separate eval source\n\n# V4.2: Stratified eval set specification (Change 1)\nEVAL_SAMPLES_PER_TASK = {\n \"extraction\": 20,\n \"sql_qa\": 15,\n \"insights\": 15,\n \"push\": 15,\n}\nEVAL_TOTAL = sum(EVAL_SAMPLES_PER_TASK.values()) # 65\n\n# ββ GRPO Hyperparameters βββββββββββββββββββββββββββββββββββββββββββββββββββββ\n# V4.2 CHANGES: MAX_STEPS 600β1500, EVAL_STEPS 20β50, SAVE_STEPS 50β100\n# Everything else UNCHANGED from V4.1 (validated)\nNUM_GENERATIONS = 16 # 0.5B + short completions = VRAM allows G=16\nMAX_COMPLETION_LENGTH = 512 # Instruct: no <think> overhead\nTEMPERATURE = 1.0 # Skywork-OR1: Ο=1.0 for exploration\nLEARNING_RATE = 5e-6 # V4.1: validated at 5e-6\nBETA = 0.0 # Dr. GRPO Β§3.2: Ξ²=0 optimal for rule-based rewards\nSCALE_REWARDS = False # Dr. GRPO: remove std normalization bias\nBATCH_SIZE = 2 # per-device batch size\nGRAD_ACCUM = 1 # effective batch = 2 * 1 = 2 prompts * 16 gen = 32 completions\nMAX_STEPS = 1500 # V4.2: was 600. ~2.5 full epochs with shuffling\nSAVE_STEPS = 100 # V4.2: was 50. Scaled for longer run\nEVAL_STEPS = 50 # V4.2: was 20. More frequent per-epoch boundary\nEARLY_STOPPING_PATIENCE = 15 # 15 Γ 50 = 750 steps without improvement\nEARLY_STOPPING_DELTA = 0.005\nLR_SCHEDULER_TYPE = \"constant_with_warmup\" # V4.1: validated\nWARMUP_RATIO = 0.05 # V4.1: validated (5% of 1500 = 75 warmup steps)\n\n# ββ LoRA βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ\nLORA_R = 16\nLORA_ALPHA = 32\n\n# ββ Monitoring βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ\nWANDB_PROJECT = \"tucano2-commerce\"\nEVAL_MAX_TOKENS = 512 # match training completion length\n\n# ββ Task Classification (inherited from V2/V3) ββββββββββββββββββββββββββββββ\nVALID_SENTIMENTS = {\"positive\", \"negative\", \"neutral\"}\nVALID_CATEGORIES = {\n \"delivery_delay\", \"product_quality\", \"product_not_received\",\n \"wrong_product\", \"seller_communication\", \"app_issue\",\n \"price_value\", \"other\", \"none\",\n}\nVALID_CHURN = {\"low\", \"medium\", \"high\"}\nVALID_REPEAT = {\"yes\", \"no\", \"maybe\"}\nEXTRACTION_FIELDS = [\n \"sentiment\", \"sentiment_score\", \"churn_risk\", \"delivery_issue\",\n \"product_issue\", \"seller_issue\", \"main_complaint\",\n \"complaint_category\", \"repeat_intent\", \"would_recommend\",\n]\n\n# ββ Verified Special Token IDs (from tokenizer_config.json) βββββββββββββββββ\nTOKEN_ID_BOS = 1 # <|im_start|>\nTOKEN_ID_EOS = 2 # <|im_end|>\nTOKEN_ID_PAD = 49109 # <|pad|>\nTOKEN_ID_THINK = 49116 # <think>\nTOKEN_ID_THINK_END = 49117 # </think>\n\n# ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ\n# TASK-AWARE SYSTEM PROMPTS (inherited from V3/V4)\n# ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ\n\nSYSTEM_EXTRACTION = (\n \"VocΓͺ Γ© um motor de extração de dados de e-commerce brasileiro. \"\n \"Retorne APENAS um objeto JSON vΓ‘lido, sem nenhum texto antes ou depois. \"\n \"NΓO USE blocos de cΓ³digo markdown (```json). \"\n \"O primeiro caractere da sua resposta deve ser { e o ΓΊltimo deve ser }. \"\n \"Campos nΓ£o mencionados na avaliação devem ser null β nunca invente valores. \"\n \"Sem explicação. Sem comentΓ‘rios.\"\n)\n\nSYSTEM_SQL = (\n \"VocΓͺ Γ© um assistente de IA especializado em anΓ‘lise de e-commerce brasileiro. \"\n \"VocΓͺ compreende avaliaçáes de clientes em portuguΓͺs e padrΓ΅es de comΓ©rcio brasileiro.\\n\\n\"\n \"Para consultas e anΓ‘lises de dados: apresente a resposta de forma direta \"\n \"com nΓΊmeros e dados concretos. Seja conciso.\"\n)\n\nSYSTEM_INSIGHTS = (\n \"VocΓͺ Γ© um assistente de IA especializado em anΓ‘lise de e-commerce brasileiro. \"\n \"VocΓͺ compreende avaliaçáes de clientes em portuguΓͺs e padrΓ΅es de comΓ©rcio brasileiro.\\n\\n\"\n \"Para anΓ‘lises estratΓ©gicas: raciocine de forma estruturada e concisa, \"\n \"focando nos pontos principais e recomendaçáes acionΓ‘veis.\"\n)\n\nSYSTEM_PUSH = (\n \"VocΓͺ Γ© um assistente de IA especializado em anΓ‘lise de e-commerce brasileiro. \"\n \"VocΓͺ compreende avaliaçáes de clientes em portuguΓͺs e padrΓ΅es de comΓ©rcio brasileiro.\\n\\n\"\n \"Para notificaçáes push: seja direto e criativo. \"\n \"A notificação deve ter no mΓ‘ximo 120 caracteres. \"\n \"Responda diretamente.\"\n)\n\nSYSTEM_PT = (\n \"VocΓͺ Γ© um assistente de IA especializado em anΓ‘lise de e-commerce brasileiro. \"\n \"VocΓͺ compreende avaliaçáes de clientes em portuguΓͺs e padrΓ΅es de comΓ©rcio brasileiro.\"\n)\n\ndef get_system_prompt(task_type: str) -> str:\n return {\n \"extraction\": SYSTEM_EXTRACTION,\n \"sql_qa\": SYSTEM_SQL,\n \"insights\": SYSTEM_INSIGHTS,\n \"push\": SYSTEM_PUSH,\n }.get(task_type, SYSTEM_PT)\n\ndef inject_task_system_prompt(msgs, task_type):\n new_msgs = []\n system_prompt = get_system_prompt(task_type)\n has_system = False\n for m in msgs:\n if m[\"role\"] == \"system\":\n new_msgs.append({\"role\": \"system\", \"content\": system_prompt})\n has_system = True\n else:\n new_msgs.append(m)\n if not has_system:\n new_msgs.insert(0, {\"role\": \"system\", \"content\": system_prompt})\n return new_msgs\n\nprint(\"β Task-aware system prompts defined\")\n\nprint(\"β Config loaded\")\nprint(f\" Model: {MODEL_ID}\")\nprint(f\" Seed: {CURRENT_SEED} (run {SEEDS.index(CURRENT_SEED)+1}/{len(SEEDS)})\")\nprint(f\" G={NUM_GENERATIONS}, max_comp={MAX_COMPLETION_LENGTH}, temp={TEMPERATURE}\")\nprint(f\" LR={LEARNING_RATE}, Ξ²={BETA}, scale_rewards={SCALE_REWARDS}\")\nprint(f\" LR schedule: {LR_SCHEDULER_TYPE}, warmup={WARMUP_RATIO}\")\nprint(f\" LoRA r={LORA_R}, Ξ±={LORA_ALPHA}\")\nprint(f\" Max steps: {MAX_STEPS} (~{MAX_STEPS * BATCH_SIZE / 1480:.1f} epochs)\")\nprint(f\" Eval: {EVAL_TOTAL} stratified samples, every {EVAL_STEPS} steps\")\nprint(f\" Save every {SAVE_STEPS} steps\")"
|
| 43 |
+
},
|
| 44 |
+
{
|
| 45 |
+
"cell_type": "markdown",
|
| 46 |
+
"metadata": {},
|
| 47 |
+
"source": "---\n\n## Cell 4: Load Model + Apply Critical Overrides\n\n**Gate:** Model loaded, `use_cache=True`, `repetition_penalty=1.0`, `temperature=1.0`.\n\nUnchanged from V4.1."
|
| 48 |
+
},
|
| 49 |
+
{
|
| 50 |
+
"cell_type": "code",
|
| 51 |
+
"execution_count": null,
|
| 52 |
+
"metadata": {},
|
| 53 |
+
"outputs": [],
|
| 54 |
+
"source": "from unsloth import FastLanguageModel\n\nprint(\"Loading model...\")\nmodel, tokenizer = FastLanguageModel.from_pretrained(\n model_name=MODEL_ID,\n max_seq_length=MAX_SEQ_LENGTH,\n load_in_4bit=True,\n dtype=None,\n)\n\nfrom peft import LoraConfig, get_peft_model\n\nlora_config = LoraConfig(\n r=LORA_R,\n lora_alpha=LORA_ALPHA,\n target_modules=[\"q_proj\", \"k_proj\", \"v_proj\", \"o_proj\",\n \"gate_proj\", \"up_proj\", \"down_proj\"],\n lora_dropout=0,\n bias=\"none\",\n task_type=\"CAUSAL_LM\",\n)\nmodel = get_peft_model(model, lora_config)\nmodel.print_trainable_parameters()\n\n# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ\n# CRITICAL OVERRIDES\n# βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ\n\nmodel.config.use_cache = True\nmodel.generation_config.use_cache = True\nmodel.generation_config.temperature = TEMPERATURE\nmodel.generation_config.repetition_penalty = 1.0\nmodel.generation_config.do_sample = True\nmodel.generation_config.top_k = 0\nmodel.generation_config.top_p = 1.0\nmodel.generation_config.max_length = None\n\nif tokenizer.pad_token is None:\n tokenizer.pad_token = tokenizer.eos_token\n\nprint(f\"β Model loaded on {model.device}\")\nprint(f\" use_cache: {model.config.use_cache}\")\nprint(f\" temperature: {model.generation_config.temperature}\")\nprint(f\" repetition_penalty: {model.generation_config.repetition_penalty}\")\nprint(f\" top_k: {model.generation_config.top_k}\")\nprint(f\" Params: {sum(p.numel() for p in model.parameters()) / 1e6:.0f}M\")\n\ntry:\n lm_ptr = model.base_model.model.lm_head.weight.data_ptr()\n embed_ptr = model.base_model.model.model.embed_tokens.weight.data_ptr()\n tied = lm_ptr == embed_ptr\n print(f\" Tied embeddings intact: {tied}\")\n if not tied:\n print(\" β οΈ WARNING: Tied embeddings broken after LoRA patching.\")\nexcept AttributeError as e:\n print(f\" β οΈ Could not check tied embeddings: {e}\")"
|
| 55 |
+
},
|
| 56 |
+
{
|
| 57 |
+
"cell_type": "markdown",
|
| 58 |
+
"metadata": {},
|
| 59 |
+
"source": "---\n\n## Cell 5: Token ID Verification\n\n**Gate:** All token IDs match. Unchanged from V4.1."
|
| 60 |
+
},
|
| 61 |
+
{
|
| 62 |
+
"cell_type": "code",
|
| 63 |
+
"execution_count": null,
|
| 64 |
+
"metadata": {},
|
| 65 |
+
"outputs": [],
|
| 66 |
+
"source": "tok_tests = {\n \"<|im_start|>\": TOKEN_ID_BOS,\n \"<|im_end|>\": TOKEN_ID_EOS,\n \"<|pad|>\": TOKEN_ID_PAD,\n \"<think>\": TOKEN_ID_THINK,\n \"</think>\": TOKEN_ID_THINK_END,\n}\n\nall_pass = True\nfor text, expected_id in tok_tests.items():\n ids = tokenizer.encode(text, add_special_tokens=False)\n actual_id = ids[0] if len(ids) == 1 else ids\n match = (len(ids) == 1 and ids[0] == expected_id)\n status = \"β\" if match else \"β\"\n print(f\" {status} '{text}' β expected {expected_id}, got {actual_id}\")\n if not match:\n all_pass = False\n\nassert all_pass, \"Token ID mismatch detected. Update constants in Cell 3 before proceeding.\"\nprint(\"\\nβ All token IDs verified\")\n\nassert tokenizer.eos_token_id == TOKEN_ID_EOS, f\"eos_token_id mismatch: {tokenizer.eos_token_id}\"\nprint(f\"β eos_token_id = {tokenizer.eos_token_id}\")"
|
| 67 |
+
},
|
| 68 |
+
{
|
| 69 |
+
"cell_type": "markdown",
|
| 70 |
+
"metadata": {},
|
| 71 |
+
"source": "---\n\n## Cell 6: KV Cache Diagnostic\n\n**Gate:** Ratio < 3Γ β KV cache OK. Unchanged from V4.1."
|
| 72 |
+
},
|
| 73 |
+
{
|
| 74 |
+
"cell_type": "code",
|
| 75 |
+
"execution_count": null,
|
| 76 |
+
"metadata": {},
|
| 77 |
+
"outputs": [],
|
| 78 |
+
"source": "FastLanguageModel.for_inference(model)\n\n_kv_msgs = [{\"role\": \"user\", \"content\": \"Qual a categoria de reclamação mais frequente?\"}]\n_kv_text = tokenizer.apply_chat_template(_kv_msgs, tokenize=False, add_generation_prompt=True)\n_kv_inputs = tokenizer(_kv_text, return_tensors=\"pt\").to(model.device)\n\n_token_times, _past, _generated = [], None, _kv_inputs[\"input_ids\"]\nwith torch.no_grad():\n for _step in range(50):\n _t0 = time.time()\n seq_len = _generated.shape[1]\n if _past is None:\n _position_ids = torch.arange(seq_len, dtype=torch.long, device=model.device).unsqueeze(0)\n else:\n _position_ids = torch.tensor([[seq_len - 1]], dtype=torch.long, device=model.device)\n _out = model(\n input_ids=_generated[:, -1:] if _past else _generated,\n position_ids=_position_ids,\n attention_mask=torch.ones(1, seq_len, device=model.device),\n past_key_values=_past,\n use_cache=True,\n return_dict=True,\n )\n _past = _out.past_key_values\n _next = _out.logits[:, -1, :].argmax(dim=-1, keepdim=True)\n _generated = torch.cat([_generated, _next], dim=1)\n _token_times.append(time.time() - _t0)\n\n_ratio = sum(_token_times[45:]) / max(sum(_token_times[:5]), 1e-9)\nprint(f\"First 5 tok: {[f'{t*1000:.0f}ms' for t in _token_times[:5]]}\")\nprint(f\"Last 5 tok: {[f'{t*1000:.0f}ms' for t in _token_times[45:]]}\")\nprint(f\"Ratio last/first: {_ratio:.1f}x\")\nassert _ratio < 5, f\"KV cache BROKEN (ratio {_ratio:.1f}Γ). Check model.config.use_cache.\"\nprint(\"β KV cache working correctly\")\n\ndel _past, _generated, _kv_inputs, _token_times, _out\ngc.collect()\ntorch.cuda.empty_cache()"
|
| 79 |
+
},
|
| 80 |
+
{
|
| 81 |
+
"cell_type": "markdown",
|
| 82 |
+
"metadata": {},
|
| 83 |
+
"source": "---\n\n## Cell 7: Reward Functions V2\n\n**V4.2 changes (Change 3 + Change 5):**\n\n### SQL Reward Overhaul (Change 3)\n- **Tier 1 (0.30):** SQL structure detected β requires β₯3 SQL keywords (SELECT, FROM, WHERE, etc.)\n- **Tier 2 (0.25):** Answer has BOTH query AND explanation (not just domain vocabulary)\n- **Tier 3 (0.25):** Numerical specificity (concrete data in answer)\n- **Tier 4 (0.20):** Portuguese business domain coherence\n\n### GDPO Per-Component Normalization (Change 5)\n- `commerce_reward_fn` now returns both per-sample rewards AND per-component rewards\n- `gdpo_normalize()` normalizes each component independently before summing\n- Preserves ~4Γ more distinct advantage groups (GDPO Β§3.1)\n\n### Dynamic Task Weights (Change 6)\n- `_task_weights` dict tracks per-task sampling weights\n- Updated by `update_task_weights()` after each eval based on improvement rate\n- MT-GRPO IWU: stagnating tasks get upweighted, improving tasks get downweighted"
|
| 84 |
+
},
|
| 85 |
+
{
|
| 86 |
+
"cell_type": "code",
|
| 87 |
+
"execution_count": null,
|
| 88 |
+
"metadata": {},
|
| 89 |
+
"outputs": [],
|
| 90 |
+
"source": "import json_repair # V4.1: robust JSON parser for LLM output\n\n\ndef strip_think(text: str) -> str:\n \"\"\"Remove <think>...</think> block, return the answer portion.\"\"\"\n return re.sub(r\"<think>.*?</think>\", \"\", text, flags=re.DOTALL).strip()\n\n\ndef has_think_block(text: str) -> bool:\n return bool(re.search(r\"<think>.+</think>\", text, flags=re.DOTALL))\n\n\ndef _classify_task_type(prompt_text: str) -> str:\n p = prompt_text.lower()\n if \"retorne um objeto json\" in p or \"extraia dados\" in p or \"json\" in p:\n return \"extraction\"\n elif \"notificação push\" in p or \"notificação de reengajamento\" in p:\n return \"push\"\n elif \"perfil do cliente\" in p or \"retenção\" in p or \"anΓ‘lise\" in p or \"insight\" in p:\n return \"insights\"\n else:\n return \"sql_qa\"\n\n\n# ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ\n# V4.1: ROBUST JSON PARSER (unchanged)\n# ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ\n\ndef _normalize_pt_decimals(s: str) -> str:\n \"\"\"Convert PT-BR decimals (4,5) to JSON-valid (4.5), only outside quoted strings.\"\"\"\n result, in_string, escape_next = [], False, False\n i = 0\n while i < len(s):\n c = s[i]\n if escape_next:\n result.append(c); escape_next = False; i += 1; continue\n if c == '\\\\' and in_string:\n result.append(c); escape_next = True; i += 1; continue\n if c == '\"':\n in_string = not in_string; result.append(c); i += 1; continue\n if not in_string:\n m = re.match(r'(\\d+),(\\d+)', s[i:])\n if m:\n result.append(m.group(1) + '.' + m.group(2))\n i += len(m.group(0)); continue\n result.append(c); i += 1\n return ''.join(result)\n\n\ndef _extract_json(text: str) -> dict | None:\n \"\"\"Robust JSON extraction for Portuguese LLM output.\"\"\"\n stripped = re.sub(r'^```(?:json)?\\s*|\\s*```$', '', text.strip(), flags=re.MULTILINE).strip()\n for attempt in [stripped, _normalize_pt_decimals(stripped)]:\n try:\n result = json.loads(attempt)\n if isinstance(result, dict):\n return result\n except (json.JSONDecodeError, TypeError):\n pass\n normalized = _normalize_pt_decimals(stripped)\n try:\n result = json_repair.repair_json(normalized, return_objects=True)\n if isinstance(result, dict):\n return result\n except Exception:\n pass\n try:\n result = json_repair.repair_json(stripped, return_objects=True)\n if isinstance(result, dict):\n return result\n except Exception:\n pass\n return None\n\n\ndef reward_extraction(completion: str) -> float:\n \"\"\"Continuous reward for extraction tasks (max 1.0). Unchanged from V4.1.\"\"\"\n answer = strip_think(completion)\n data = _extract_json(answer)\n\n if data is None:\n if \"{\" in answer and \"}\" in answer:\n return 0.05\n return 0.0\n\n if not isinstance(data, dict):\n return 0.1\n\n score = 0.3 # valid JSON object\n\n # Schema completeness (0.3 total)\n present = sum(1 for f in EXTRACTION_FIELDS if f in data)\n score += 0.3 * (present / len(EXTRACTION_FIELDS))\n\n # Value validity (0.4 total)\n checks_passed = 0\n checks_total = 0\n\n for field, validator in [\n (\"sentiment\", lambda v: isinstance(v, str) and v in VALID_SENTIMENTS),\n (\"complaint_category\", lambda v: isinstance(v, str) and v in VALID_CATEGORIES),\n (\"churn_risk\", lambda v: isinstance(v, str) and v in VALID_CHURN),\n (\"repeat_intent\", lambda v: isinstance(v, str) and v in VALID_REPEAT),\n (\"sentiment_score\", lambda v: isinstance(v, (int, float)) and 1 <= v <= 5),\n ]:\n checks_total += 1\n if field in data and validator(data[field]):\n checks_passed += 1\n\n for bool_field in (\"delivery_issue\", \"product_issue\", \"seller_issue\", \"would_recommend\"):\n checks_total += 1\n if bool_field in data and isinstance(data[bool_field], bool):\n checks_passed += 1\n\n if checks_total > 0:\n score += 0.4 * (checks_passed / checks_total)\n\n return min(score, 1.0)\n\n\n# ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ\n# V4.2: SQL REWARD V2 β Validation-aware (Change 3)\n# Replaces heuristic vocabulary matching with structural analysis.\n# Expected: distinguishes \"mentions SQL keywords\" from \"produces correct answer\"\n# ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ\n\ndef reward_sql_qa(completion: str) -> float:\n \"\"\"V4.2: Validation-aware SQL Q&A reward (max 1.0).\n \n Tier 1 (0.30): SQL structure detected (β₯3 keywords or code block)\n Tier 2 (0.25): Answer has both query and explanation\n Tier 3 (0.25): Numerical specificity (concrete data)\n Tier 4 (0.20): Portuguese business domain coherence\n \"\"\"\n answer = strip_think(completion)\n if not answer.strip():\n return 0.0\n\n score = 0.0\n\n # Tier 1 (0.30): SQL structure detected\n sql_keywords = [\"SELECT\", \"FROM\", \"WHERE\", \"GROUP BY\", \"ORDER BY\",\n \"JOIN\", \"HAVING\", \"COUNT\", \"AVG\", \"SUM\"]\n sql_found = sum(1 for kw in sql_keywords if kw in answer.upper())\n if sql_found >= 3:\n score += 0.30\n elif sql_found >= 1:\n score += 0.15\n\n # Tier 2 (0.25): Answer has both query AND explanation\n has_query = bool(re.search(r\"```sql|SELECT.{5,}FROM\", answer, re.IGNORECASE | re.DOTALL))\n has_answer = any(w in answer.lower() for w in [\"resultado\", \"total\", \"mΓ©dia\", \"mostra\", \"portanto\"])\n if has_query and has_answer:\n score += 0.25\n elif has_query or has_answer:\n score += 0.12\n\n # Tier 3 (0.25): Numerical specificity\n numbers = re.findall(r\"\\d+(?:[.,]\\d+)?(?:\\s*%)?\", answer)\n score += min(0.25, 0.05 * len(numbers))\n\n # Tier 4 (0.20): Portuguese business domain coherence\n pt_domain = [\"pedidos\", \"clientes\", \"vendedores\", \"produtos\", \"avaliação\",\n \"entrega\", \"reclamação\", \"satisfação\", \"categoria\", \"perΓodo\"]\n score += min(0.20, 0.04 * sum(1 for w in pt_domain if w in answer.lower()))\n\n return min(score, 1.0)\n\n\ndef reward_insights(completion: str) -> float:\n \"\"\"Continuous reward for insights (max 1.0). Unchanged from V4.1.\"\"\"\n answer = strip_think(completion)\n if not answer.strip():\n return 0.0\n\n score = 0.0\n\n action_words = [\"recomend\", \"implement\", \"melhor\", \"reduzir\", \"aumentar\",\n \"priorizar\", \"investir\", \"otimizar\", \"estratΓ©gi\", \"ação\"]\n matches = sum(1 for w in action_words if w in answer.lower())\n score += min(0.4, 0.08 * matches)\n\n length = len(answer)\n if 100 <= length <= 800:\n score += 0.3\n elif length > 0:\n score += 0.3 * max(0, 1 - abs(length - 450) / 450)\n\n structure_marks = len(re.findall(r\"^[-β’*]\\s|^\\d+[.)]\\s|^#{1,3}\\s\", answer, re.MULTILINE))\n score += min(0.2, 0.04 * structure_marks)\n\n if any(w in answer.lower() for w in [\"cliente\", \"produto\", \"serviΓ§o\", \"empresa\"]):\n score += 0.1\n\n return min(score, 1.0)\n\n\ndef reward_push(completion: str) -> float:\n \"\"\"Continuous reward for push notifications (max 1.0). Unchanged from V4.1.\"\"\"\n answer = strip_think(completion).strip()\n if not answer:\n return 0.0\n\n length = len(answer)\n if length <= 120:\n length_score = 0.5\n else:\n length_score = 0.5 * max(0, 1 - (length - 120) / 120)\n\n pt_markers = re.findall(r\"[ãçéΓͺΓ³ΓΊΓ’Γ΅]|vocΓͺ|para|como|seu|sua|oferta|desconto|produto\",\n answer, re.IGNORECASE)\n lang_score = min(0.3, 0.03 * len(pt_markers))\n\n generic = [\"olΓ‘\", \"obrigado pela compra\", \"agradecemos\"]\n is_generic = any(g in answer.lower() for g in generic)\n creativity_score = 0.0 if is_generic else 0.2\n\n return min(length_score + lang_score + creativity_score, 1.0)\n\n\n# ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ\n# V4.2: GDPO PER-COMPONENT NORMALIZATION (Change 5)\n# Normalize each reward component independently before aggregation.\n# GDPO (2601.05242) Β§3.1: preserves ~4Γ more distinct advantage groups.\n# ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ\n\ndef gdpo_normalize(component_rewards: dict) -> list:\n \"\"\"Per-component normalization before aggregation (GDPO 2601.05242 Β§3.1).\n \n Args:\n component_rewards: {task_name: [reward_per_sample, ...]} for each component\n \n Returns:\n List of normalized summed rewards, one per sample.\n \"\"\"\n normalized = {}\n for task, rewards in component_rewards.items():\n rewards_t = torch.tensor(rewards, dtype=torch.float32)\n std = rewards_t.std()\n if std > 1e-8:\n normalized[task] = ((rewards_t - rewards_t.mean()) / std).tolist()\n else:\n normalized[task] = [0.0] * len(rewards) # zero-variance group\n # Sum normalized components per sample\n n = len(next(iter(normalized.values())))\n return [sum(normalized[t][i] for t in normalized) for i in range(n)]\n\n\n# ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ\n# V4.2: DYNAMIC TASK WEIGHTING β MT-GRPO IWU (Change 6)\n# Track per-task reward improvement rates, upweight stagnating tasks.\n# MT-GRPO (2602.05547) Β§3.2: prevents easy-task collapse.\n# ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ\n\n_task_weights = {\n \"extraction\": 0.25,\n \"sql_qa\": 0.25,\n \"insights\": 0.25,\n \"push\": 0.25,\n}\n_task_reward_history = {t: [] for t in _task_weights}\n\ndef update_task_weights(step: int, per_task_rewards: dict, update_interval: int = 50):\n \"\"\"MT-GRPO IWU: update task sampling weights based on improvement rate.\n \n Args:\n step: Current training step\n per_task_rewards: {task: mean_reward} from latest eval\n update_interval: Only update every N steps\n \"\"\"\n global _task_weights\n if step % update_interval != 0 or step == 0:\n return\n \n for task, reward in per_task_rewards.items():\n if task not in _task_reward_history:\n continue\n _task_reward_history[task].append(reward)\n if len(_task_reward_history[task]) >= 2:\n improvement = _task_reward_history[task][-1] - _task_reward_history[task][-2]\n if improvement < 0.01: # stagnating\n _task_weights[task] = min(0.60, _task_weights[task] * 1.3)\n elif improvement > 0.05: # improving fast\n _task_weights[task] = max(0.10, _task_weights[task] * 0.85)\n \n # Normalize to sum to 1\n total = sum(_task_weights.values())\n _task_weights = {t: w / total for t, w in _task_weights.items()}\n\n\ndef get_task_weighted_indices(dataset, n_samples: int) -> list:\n \"\"\"Sample indices with probability proportional to task weights.\"\"\"\n task_indices = {t: [] for t in _task_weights}\n for i, record in enumerate(dataset):\n user_txt = \" \".join(m[\"content\"] for m in record[\"prompt\"] if m[\"role\"] == \"user\")\n task = _classify_task_type(user_txt)\n if task in task_indices:\n task_indices[task].append(i)\n \n sampled = []\n for task, weight in _task_weights.items():\n n = max(1, int(n_samples * weight))\n pool = task_indices.get(task, [])\n if pool:\n sampled.extend(random.sample(pool, min(n, len(pool))))\n random.shuffle(sampled)\n return sampled[:n_samples]\n\n\n# ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ\n# MASTER REWARD FUNCTION β V4.2: returns per-component rewards for GDPO\n# ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ\n\ndef commerce_reward_fn(completions, prompts, **kwargs) -> list:\n \"\"\"Master reward function: dispatches by task type.\n \n V4.2: Also stores per-component rewards in _last_component_rewards\n for GDPO normalization when called externally.\n \"\"\"\n rewards = []\n for completion, prompt in zip(completions, prompts):\n if isinstance(completion, list):\n comp_text = completion[-1][\"content\"] if completion else \"\"\n else:\n comp_text = str(completion)\n\n if isinstance(prompt, list):\n prompt_text = \" \".join(m.get(\"content\", \"\") for m in prompt)\n else:\n prompt_text = str(prompt)\n\n task = _classify_task_type(prompt_text)\n\n if task == \"extraction\":\n rewards.append(reward_extraction(comp_text))\n elif task == \"sql_qa\":\n rewards.append(reward_sql_qa(comp_text))\n elif task == \"insights\":\n rewards.append(reward_insights(comp_text))\n elif task == \"push\":\n rewards.append(reward_push(comp_text))\n else:\n r = 0.2 if comp_text.strip() else 0.0\n rewards.append(r)\n\n return rewards\n\n\ndef commerce_reward_fn_with_components(completions, prompts, **kwargs):\n \"\"\"V4.2: Reward function that also returns per-task component rewards for GDPO.\n \n Returns:\n tuple: (rewards_list, component_dict)\n component_dict: {task: [reward_for_each_sample_of_this_task]}\n \"\"\"\n rewards = []\n components = {\"extraction\": [], \"sql_qa\": [], \"insights\": [], \"push\": []}\n indices = {\"extraction\": [], \"sql_qa\": [], \"insights\": [], \"push\": []}\n \n for idx, (completion, prompt) in enumerate(zip(completions, prompts)):\n if isinstance(completion, list):\n comp_text = completion[-1][\"content\"] if completion else \"\"\n else:\n comp_text = str(completion)\n\n if isinstance(prompt, list):\n prompt_text = \" \".join(m.get(\"content\", \"\") for m in prompt)\n else:\n prompt_text = str(prompt)\n\n task = _classify_task_type(prompt_text)\n\n if task == \"extraction\":\n r = reward_extraction(comp_text)\n elif task == \"sql_qa\":\n r = reward_sql_qa(comp_text)\n elif task == \"insights\":\n r = reward_insights(comp_text)\n elif task == \"push\":\n r = reward_push(comp_text)\n else:\n r = 0.2 if comp_text.strip() else 0.0\n task = \"sql_qa\" # default bucket\n\n rewards.append(r)\n if task in components:\n components[task].append(r)\n indices[task].append(idx)\n\n return rewards, components, indices\n\n\nprint(\"β Reward functions defined (V4.2: SQL v2 + GDPO + dynamic weights)\")\nprint(f\" Task weights: {_task_weights}\")"
|
| 91 |
+
},
|
| 92 |
+
{
|
| 93 |
+
"cell_type": "markdown",
|
| 94 |
+
"metadata": {},
|
| 95 |
+
"source": "---\n\n## Cell 8: Reward Function Audit (Change 2)\n\n**V4.2 addition: 30-minute audit protocol.**\n\nGenerate 20 completions (5 per task) at temp=0.1, score them with the reward function,\nthen have the human score them 0-10. Compute Spearman Ο.\n\n**Gate:** Ο > 0.70. If below, reward function is miscalibrated β fix before training.\n\n**Why:** The V1-V4 parser bug would have been caught in 30 minutes with this protocol.\n\n### Instructions\n1. Run this cell β it generates 20 completions and scores them automatically\n2. Read each completion and assign a 0-10 human quality score\n3. Fill in the `audit_human_scores` list\n4. Re-run the assertion at the bottom\n5. If Ο < 0.70, investigate discrepancies before proceeding to training"
|
| 96 |
+
},
|
| 97 |
+
{
|
| 98 |
+
"cell_type": "code",
|
| 99 |
+
"execution_count": null,
|
| 100 |
+
"metadata": {},
|
| 101 |
+
"outputs": [],
|
| 102 |
+
"source": "from scipy.stats import spearmanr\n\nAUDIT_PROMPTS_PER_TASK = 5\n\n# ββ Collect audit prompts (5 per task) βββββββββββββββββββββββββββββββββββββββ\naudit_by_type = {\"extraction\": [], \"sql_qa\": [], \"insights\": [], \"push\": []}\nwith open(TRAIN_FILE) as f:\n for line in f:\n row = json.loads(line)\n convs = row[\"conversations\"]\n prompt_msgs = [m for m in convs if m[\"role\"] in (\"system\", \"user\")]\n if not prompt_msgs:\n continue\n user_text = \" \".join(m[\"content\"] for m in prompt_msgs if m[\"role\"] == \"user\")\n task = _classify_task_type(user_text)\n if len(audit_by_type[task]) < AUDIT_PROMPTS_PER_TASK:\n audit_by_type[task].append(prompt_msgs)\n\nprint(f\"Audit prompts collected: {', '.join(f'{k}={len(v)}' for k, v in audit_by_type.items())}\")\n\n# ββ Generate completions and score automatically βββββββββββββββββββββββββββββ\nFastLanguageModel.for_inference(model)\n\naudit_auto_scores = []\naudit_tasks = []\naudit_completions = []\n\nfor task_type in [\"extraction\", \"sql_qa\", \"insights\", \"push\"]:\n for msgs in audit_by_type[task_type]:\n msgs = inject_task_system_prompt(msgs, task_type)\n text = tokenizer.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True)\n inputs = tokenizer(text, return_tensors=\"pt\").to(model.device)\n with torch.no_grad():\n out = model.generate(\n **inputs,\n max_new_tokens=MAX_COMPLETION_LENGTH,\n temperature=0.1, # near-deterministic for audit\n do_sample=True,\n repetition_penalty=1.0,\n )\n resp = tokenizer.decode(out[0][inputs[\"input_ids\"].shape[1]:], skip_special_tokens=True)\n r = commerce_reward_fn([resp], [text])[0]\n audit_auto_scores.append(r)\n audit_tasks.append(task_type)\n audit_completions.append(resp)\n\nprint(f\"\\n{'='*80}\")\nprint(\"REWARD FUNCTION AUDIT β 20 Completions\")\nprint(f\"{'='*80}\")\nfor i, (task, auto_r, comp) in enumerate(zip(audit_tasks, audit_auto_scores, audit_completions)):\n answer = strip_think(comp)[:200]\n print(f\"\\n--- Sample {i+1}/{len(audit_auto_scores)} [{task}] auto_reward={auto_r:.3f} ---\")\n print(f\"{answer}\")\n\nprint(f\"\\n{'='*80}\")\nprint(\"INSTRUCTIONS:\")\nprint(\"1. Read each completion above\")\nprint(\"2. Assign a 0-10 quality score for each (0=garbage, 10=perfect)\")\nprint(\"3. Fill in the list below and re-run the assertion\")\nprint(f\"{'='*80}\")\n\n# ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ\n# HUMAN SCORES β fill this in after reading the completions above\n# Each score is 0-10. Order matches the samples printed above.\n# ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ\naudit_human_scores = [\n # extraction (5 scores)\n -1, -1, -1, -1, -1,\n # sql_qa (5 scores)\n -1, -1, -1, -1, -1,\n # insights (5 scores)\n -1, -1, -1, -1, -1,\n # push (5 scores)\n -1, -1, -1, -1, -1,\n]\n\n# ββ Compute correlation ββββββββββββββββββββββββββββββββββββββββββββββββββββββ\nif all(s >= 0 for s in audit_human_scores):\n # Normalize human scores to 0-1 range for comparison\n human_normalized = [s / 10.0 for s in audit_human_scores]\n rho, p_value = spearmanr(human_normalized, audit_auto_scores)\n print(f\"\\nReward function calibration: Ο={rho:.3f} (p={p_value:.4f})\")\n print(f\" Human scores (normalized): {[f'{s:.1f}' for s in human_normalized]}\")\n print(f\" Auto scores: {[f'{s:.3f}' for s in audit_auto_scores]}\")\n \n if rho > 0.70:\n print(f\" β
PASS: Ο={rho:.3f} > 0.70 β reward function is calibrated\")\n else:\n print(f\" β FAIL: Ο={rho:.3f} < 0.70 β reward function is miscalibrated\")\n print(\" β Investigate discrepancies before training. Check:\")\n print(\" 1. Is the JSON parser handling all formats correctly?\")\n print(\" 2. Are SQL reward tiers appropriate for this model's output style?\")\n print(\" 3. Are insights/push length penalties calibrated?\")\n # Show biggest discrepancies\n diffs = [(i, abs(human_normalized[i] - audit_auto_scores[i]), audit_tasks[i])\n for i in range(len(audit_human_scores))]\n diffs.sort(key=lambda x: x[1], reverse=True)\n print(f\"\\n Top 5 discrepancies:\")\n for idx, diff, task in diffs[:5]:\n print(f\" Sample {idx+1} [{task}]: human={human_normalized[idx]:.1f}, auto={audit_auto_scores[idx]:.3f}, Ξ={diff:.3f}\")\n \n assert rho > 0.70, f\"Reward function miscalibrated (Ο={rho:.3f} < 0.70). Fix before training.\"\nelse:\n print(\"\\nβ οΈ Human scores not yet filled in. Fill audit_human_scores and re-run.\")\n print(\" You can proceed to Cell 9 to build the eval set while scoring.\")\n print(\" But DO NOT proceed past Cell 11 (smoke test) without completing the audit.\")"
|
| 103 |
+
},
|
| 104 |
+
{
|
| 105 |
+
"cell_type": "markdown",
|
| 106 |
+
"metadata": {},
|
| 107 |
+
"source": "---\n\n## Cell 9: Build Stratified Eval Set (Change 1)\n\n**V4.2: 65 stratified samples** (20 extraction + 15 sql_qa + 15 insights + 15 push).\n\nSampled from `data/pairs/eval.jsonl`, saved as `data/pairs/eval_v2_stratified.jsonl`.\n**This file is fixed across all seeds.** Never resample.\n\n**Gate:** Exactly 65 samples with correct per-task counts."
|
| 108 |
+
},
|
| 109 |
+
{
|
| 110 |
+
"cell_type": "code",
|
| 111 |
+
"execution_count": null,
|
| 112 |
+
"metadata": {},
|
| 113 |
+
"outputs": [],
|
| 114 |
+
"source": "eval_v2_stratified_path = DATA_DIR / \"pairs\" / \"eval_v2_stratified.jsonl\"\n\n# ββ Check if already built (idempotent across seeds) ββββββββββββββββββββββββ\nif eval_v2_stratified_path.exists():\n existing = []\n with open(eval_v2_stratified_path) as f:\n for line in f:\n existing.append(json.loads(line))\n # Verify counts\n task_counts = {}\n for rec in existing:\n task_counts[rec[\"task_type\"]] = task_counts.get(rec[\"task_type\"], 0) + 1\n print(f\"β Stratified eval set already exists: {eval_v2_stratified_path}\")\n print(f\" Counts: {task_counts}\")\n print(f\" Total: {len(existing)}\")\n assert len(existing) == EVAL_TOTAL, f\"Expected {EVAL_TOTAL}, got {len(existing)}\"\nelse:\n # ββ Build from eval.jsonl (or train.jsonl fallback) ββββββββββββββββββββββ\n eval_source = EVAL_FILE if EVAL_FILE.exists() else TRAIN_FILE\n print(f\"Building stratified eval set from: {eval_source}\")\n \n # Collect all records by task\n eval_by_task = {t: [] for t in EVAL_SAMPLES_PER_TASK}\n with open(eval_source) as f:\n for line in f:\n row = json.loads(line)\n convs = row[\"conversations\"]\n prompt_msgs = [m for m in convs if m[\"role\"] in (\"system\", \"user\")]\n if not prompt_msgs:\n continue\n user_text = \" \".join(m[\"content\"] for m in prompt_msgs if m[\"role\"] == \"user\")\n task = _classify_task_type(user_text)\n if task in eval_by_task:\n eval_by_task[task].append({\n \"conversations\": convs,\n \"prompt_msgs\": prompt_msgs,\n \"task_type\": task,\n })\n \n print(f\" Available: {', '.join(f'{k}={len(v)}' for k, v in eval_by_task.items())}\")\n \n # Stratified sampling with FIXED seed (42 always, regardless of CURRENT_SEED)\n eval_rng = random.Random(42)\n stratified = []\n for task, target_n in EVAL_SAMPLES_PER_TASK.items():\n pool = eval_by_task[task]\n if len(pool) < target_n:\n print(f\" β οΈ {task}: only {len(pool)} available, wanted {target_n}. Using all.\")\n sampled = pool\n else:\n sampled = eval_rng.sample(pool, target_n)\n stratified.extend(sampled)\n \n # Save as JSONL\n eval_v2_stratified_path.parent.mkdir(parents=True, exist_ok=True)\n with open(eval_v2_stratified_path, \"w\") as f:\n for rec in stratified:\n f.write(json.dumps(rec, ensure_ascii=False) + \"\\n\")\n \n # Verify\n task_counts = {}\n for rec in stratified:\n task_counts[rec[\"task_type\"]] = task_counts.get(rec[\"task_type\"], 0) + 1\n \n print(f\"\\nβ Stratified eval set saved: {eval_v2_stratified_path}\")\n print(f\" Counts: {task_counts}\")\n print(f\" Total: {len(stratified)}\")\n\nprint(f\"\\nExpected: {EVAL_SAMPLES_PER_TASK} = {EVAL_TOTAL} total\")"
|
| 115 |
+
},
|
| 116 |
+
{
|
| 117 |
+
"cell_type": "markdown",
|
| 118 |
+
"metadata": {},
|
| 119 |
+
"source": "---\n\n## Cell 10: Dataset Preparation + DynamicTaskSampler Init\n\n**V4.2 changes:**\n- Eval loaded from `eval_v2_stratified.jsonl` (65 fixed samples) instead of random split\n- Train still from `train.jsonl` with task-aware system prompt injection\n- DynamicTaskSampler initialized for IWU (Change 6)\n\n**Gate:** Train has ~1,480+ prompts. Eval has exactly 65 stratified samples. All 4 task types present in both."
|
| 120 |
+
},
|
| 121 |
+
{
|
| 122 |
+
"cell_type": "code",
|
| 123 |
+
"execution_count": null,
|
| 124 |
+
"metadata": {},
|
| 125 |
+
"outputs": [],
|
| 126 |
+
"source": "from datasets import Dataset\n\ndef prepare_datasets_v42(train_file, eval_stratified_file, seed=CURRENT_SEED):\n \"\"\"V4.2: Load train from JSONL + eval from stratified file.\"\"\"\n rng = random.Random(seed)\n\n # ββ Train set ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ\n train_records = []\n with open(train_file) as f:\n for line in f:\n row = json.loads(line)\n convs = row[\"conversations\"]\n prompt_msgs = [m for m in convs if m[\"role\"] in (\"system\", \"user\")]\n if prompt_msgs:\n train_records.append(prompt_msgs)\n rng.shuffle(train_records)\n \n # Inject task-aware system prompts\n for i, msgs in enumerate(train_records):\n user_text = \" \".join(m[\"content\"] for m in msgs if m[\"role\"] == \"user\")\n task = _classify_task_type(user_text)\n train_records[i] = inject_task_system_prompt(msgs, task)\n \n # ββ Eval set (V4.2: from stratified file, fixed across seeds) ββββββββββββ\n eval_records = []\n with open(eval_stratified_file) as f:\n for line in f:\n rec = json.loads(line)\n prompt_msgs = rec[\"prompt_msgs\"]\n user_text = \" \".join(m[\"content\"] for m in prompt_msgs if m[\"role\"] == \"user\")\n task = _classify_task_type(user_text)\n eval_records.append(inject_task_system_prompt(prompt_msgs, task))\n \n print(\" β Task-aware system prompts injected\")\n \n # Print distributions\n for label, records in [(\"train\", train_records), (\"eval\", eval_records)]:\n dist = {}\n for msgs in records:\n user_text = \" \".join(m[\"content\"] for m in msgs if m[\"role\"] == \"user\")\n task = _classify_task_type(user_text)\n dist[task] = dist.get(task, 0) + 1\n print(f\" {label}: {len(records)} prompts β {dist}\")\n \n train_ds = Dataset.from_list([{\"prompt\": msgs} for msgs in train_records])\n eval_ds = Dataset.from_list([{\"prompt\": msgs} for msgs in eval_records])\n return train_ds, eval_ds\n\ntrain_dataset, eval_dataset = prepare_datasets_v42(TRAIN_FILE, eval_v2_stratified_path)\nprint(f\"\\nβ Datasets: train={len(train_dataset)}, eval={len(eval_dataset)}\")\nassert len(eval_dataset) == EVAL_TOTAL, f\"Expected {EVAL_TOTAL} eval samples, got {len(eval_dataset)}\"\nprint(f\"β Eval is stratified: {EVAL_TOTAL} samples (fixed across seeds)\")"
|
| 127 |
+
},
|
| 128 |
+
{
|
| 129 |
+
"cell_type": "markdown",
|
| 130 |
+
"metadata": {},
|
| 131 |
+
"source": "---\n\n## Cell 11: Smoke Test (1 Step)\n\n**Gate:** No OOM. Peak VRAM < 20GB. Step time < 180s."
|
| 132 |
+
},
|
| 133 |
+
{
|
| 134 |
+
"cell_type": "code",
|
| 135 |
+
"execution_count": null,
|
| 136 |
+
"metadata": {},
|
| 137 |
+
"outputs": [],
|
| 138 |
+
"source": "from trl import GRPOConfig, GRPOTrainer\n\nFastLanguageModel.for_training(model)\n\nsmoke_config = GRPOConfig(\n output_dir=str(CHECKPOINT_DIR / \"smoke\"),\n num_generations=NUM_GENERATIONS,\n scale_rewards=SCALE_REWARDS,\n max_completion_length=MAX_COMPLETION_LENGTH,\n max_steps=1,\n temperature=TEMPERATURE,\n beta=BETA,\n per_device_train_batch_size=BATCH_SIZE,\n gradient_accumulation_steps=1,\n learning_rate=LEARNING_RATE,\n lr_scheduler_type=LR_SCHEDULER_TYPE,\n warmup_ratio=WARMUP_RATIO,\n fp16=False,\n bf16=True,\n logging_steps=1,\n save_steps=999,\n report_to=\"none\",\n max_prompt_length=MAX_SEQ_LENGTH // 2,\n seed=CURRENT_SEED,\n remove_unused_columns=False,\n)\n\n\nclass UnslothGRPOTrainer(GRPOTrainer):\n def _generate(self, prompts, images):\n FastLanguageModel.for_inference(self.model)\n try:\n result = super()._generate(prompts, images)\n finally:\n FastLanguageModel.for_training(self.model)\n return result\n\n\nsmoke_trainer = UnslothGRPOTrainer(\n model=model,\n reward_funcs=commerce_reward_fn,\n args=smoke_config,\n train_dataset=train_dataset,\n processing_class=tokenizer,\n)\n\nt0 = time.time()\nsmoke_trainer.train()\nstep_time = time.time() - t0\n\npeak_vram = torch.cuda.max_memory_allocated() / 1e9\nprint(f\"\\nβ Smoke test passed!\")\nprint(f\" Step time: {step_time:.0f}s\")\nprint(f\" Peak VRAM: {peak_vram:.1f}GB / {torch.cuda.get_device_properties(0).total_mem / 1e9:.1f}GB\")\nprint(f\" Estimated full run ({MAX_STEPS} steps): {step_time * MAX_STEPS / 3600:.1f}h\")\n\ndel smoke_trainer\ngc.collect(); torch.cuda.empty_cache()"
|
| 139 |
+
},
|
| 140 |
+
{
|
| 141 |
+
"cell_type": "markdown",
|
| 142 |
+
"metadata": {},
|
| 143 |
+
"source": "---\n\n## Cell 12: Probe Run (10 Steps)\n\n**V4.2:** Uses `CURRENT_SEED` for reproducibility. No hard clip_ratio gate (expected=0 for LoRA)."
|
| 144 |
+
},
|
| 145 |
+
{
|
| 146 |
+
"cell_type": "code",
|
| 147 |
+
"execution_count": null,
|
| 148 |
+
"metadata": {},
|
| 149 |
+
"outputs": [],
|
| 150 |
+
"source": "FastLanguageModel.for_training(model)\n\nprobe_config = GRPOConfig(\n output_dir=str(CHECKPOINT_DIR / \"probe\"),\n num_generations=NUM_GENERATIONS,\n scale_rewards=SCALE_REWARDS,\n max_completion_length=MAX_COMPLETION_LENGTH,\n max_steps=10,\n temperature=TEMPERATURE,\n beta=BETA,\n num_train_epochs=1,\n per_device_train_batch_size=BATCH_SIZE,\n gradient_accumulation_steps=GRAD_ACCUM,\n learning_rate=LEARNING_RATE,\n lr_scheduler_type=LR_SCHEDULER_TYPE,\n warmup_ratio=WARMUP_RATIO,\n fp16=False,\n bf16=True,\n logging_steps=1,\n save_steps=999,\n report_to=\"none\",\n max_prompt_length=MAX_SEQ_LENGTH // 2,\n seed=CURRENT_SEED,\n remove_unused_columns=False,\n)\n\nprobe_trainer = UnslothGRPOTrainer(\n model=model,\n reward_funcs=commerce_reward_fn,\n args=probe_config,\n train_dataset=train_dataset,\n processing_class=tokenizer,\n)\n\nt0 = time.time()\nresult = probe_trainer.train()\nelapsed = time.time() - t0\n\n# ββ Extract metrics from log history βββββββββββββββββββββββββββββββββββββββββ\nrewards = []\ngrad_norms = []\nzero_stds = []\nfor entry in probe_trainer.state.log_history:\n if \"train/reward\" in entry:\n rewards.append(entry[\"train/reward\"])\n if \"train/grad_norm\" in entry:\n grad_norms.append(entry[\"train/grad_norm\"])\n if \"train/frac_reward_zero_std\" in entry:\n zero_stds.append(entry[\"train/frac_reward_zero_std\"])\n\nprint(f\"\\n{'='*60}\")\nprint(f\"PROBE RESULTS ({elapsed:.0f}s, {elapsed/10:.0f}s/step)\")\nprint(f\" Rewards: {[f'{r:.3f}' for r in rewards]}\")\nprint(f\" Grad norms: {[f'{g:.4f}' for g in grad_norms]}\")\nprint(f\" Zero-std: {[f'{z:.2f}' for z in zero_stds]}\")\nprint(f\" Train loss: {result.training_loss:.4f}\")\nprint(f\"{'='*60}\")\n\nif rewards and max(rewards) > 0:\n print(\"β Model generates scoreable output\")\nelse:\n print(\"β WARNING: All rewards are 0. Check reward functions.\")\n\nif grad_norms and max(grad_norms) > 0:\n print(\"β Gradients are flowing\")\nelse:\n print(\"β WARNING: All grad_norms are 0. Check model/LoRA setup.\")\n\nif zero_stds and max(zero_stds) < 0.5:\n print(\"β Batches have reward variance (GRPO has signal)\")\nelse:\n print(\"β οΈ WARNING: High frac_reward_zero_std. Consider increasing G.\")\n\nprint(\"\\nβ Proceed to full training (Cell 13)\")\n\ndel probe_trainer\ngc.collect(); torch.cuda.empty_cache()"
|
| 151 |
+
},
|
| 152 |
+
{
|
| 153 |
+
"cell_type": "markdown",
|
| 154 |
+
"metadata": {},
|
| 155 |
+
"source": "---\n\n## Cell 13: W&B Init + Full Training (1,500 Steps)\n\n**V4.2 changes:**\n- **MAX_STEPS=1,500** (multi-epoch, ~2.5Γ full dataset) (Change 4)\n- **EvalRewardCallback v2:** 65 stratified samples, per-task 95% CIs, GDPO normalization logging, dynamic task weight updates, **best checkpoint saving** (Changes 1, 5, 6, 8)\n- **`SAVE_STEPS=100`, `EVAL_STEPS=50`** (scaled for longer run)\n- **Seed in W&B config** for multi-seed tracking (Change 7)\n- **Best checkpoint saved explicitly** when eval improves (Change 8)"
|
| 156 |
+
},
|
| 157 |
+
{
|
| 158 |
+
"cell_type": "code",
|
| 159 |
+
"execution_count": null,
|
| 160 |
+
"metadata": {},
|
| 161 |
+
"outputs": [],
|
| 162 |
+
"source": "import wandb\nfrom transformers import TrainerCallback\n\nwandb.login()\nwandb.init(\n project=WANDB_PROJECT,\n name=f\"grpo-v4.2-instruct-0.5B-seed{CURRENT_SEED}-{time.strftime('%Y%m%d-%H%M')}\",\n config={\n \"model_id\": MODEL_ID,\n \"version\": \"v4.2\",\n \"seed\": CURRENT_SEED,\n \"seeds_planned\": SEEDS,\n \"num_generations\": NUM_GENERATIONS,\n \"max_completion_length\": MAX_COMPLETION_LENGTH,\n \"temperature\": TEMPERATURE,\n \"learning_rate\": LEARNING_RATE,\n \"lr_scheduler_type\": LR_SCHEDULER_TYPE,\n \"warmup_ratio\": WARMUP_RATIO,\n \"beta\": BETA,\n \"scale_rewards\": SCALE_REWARDS,\n \"batch_size\": BATCH_SIZE,\n \"grad_accum\": GRAD_ACCUM,\n \"max_steps\": MAX_STEPS,\n \"lora_r\": LORA_R,\n \"lora_alpha\": LORA_ALPHA,\n \"train_prompts\": len(train_dataset),\n \"eval_prompts\": len(eval_dataset),\n \"eval_stratified\": True,\n \"eval_per_task\": EVAL_SAMPLES_PER_TASK,\n \"repetition_penalty_override\": 1.0,\n \"json_parser\": \"json-repair + PT-BR decimal normalizer\",\n \"sql_reward\": \"v2 (validation-aware, 4-tier)\",\n \"gdpo_normalization\": True,\n \"dynamic_task_weighting\": \"MT-GRPO IWU\",\n \"changes_from_v41\": \"stratified eval 65, reward audit, SQL v2, 1500 steps, GDPO, IWU, 3 seeds, best ckpt\",\n },\n)\nprint(f\"β W&B run: {wandb.run.url}\")\n\n\n# ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ\n# V4.2: EvalRewardCallback v2\n# - Uses 65 stratified eval samples (Change 1)\n# - Reports per-task means with 95% CIs (Change 1)\n# - Runs GDPO normalization and logs component stats (Change 5)\n# - Updates dynamic task weights via IWU (Change 6)\n# - Saves best checkpoint explicitly (Change 8)\n# ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ\n\nclass EvalRewardCallbackV2(TrainerCallback):\n def __init__(self, eval_records, reward_fn, patience, delta):\n self.eval_records = eval_records\n self.reward_fn = reward_fn\n self.patience = patience\n self.delta = delta\n self.best_reward = -float(\"inf\")\n self.best_step = 0\n self.no_improve_count = 0\n\n def on_step_end(self, args, state, control, model=None, processing_class=None, **kwargs):\n if state.global_step == 0 or state.global_step % EVAL_STEPS != 0:\n return control\n\n tokenizer_local = processing_class\n if tokenizer_local is None:\n print(\"[EvalRewardCallback] WARNING: tokenizer is None, skipping eval\")\n return control\n\n mean_reward, per_task, per_task_all = self._run_eval(model, tokenizer_local, args)\n improved = mean_reward > self.best_reward + self.delta\n\n # ββ Per-task 95% CIs (Change 1) ββββββββββββββββββββββββββββββββββββββ\n log_data = {\n \"eval/mean_reward\": mean_reward,\n \"eval/best_reward\": max(self.best_reward, mean_reward),\n \"eval/no_improve_count\": self.no_improve_count,\n }\n \n ci_strs = []\n for task_name, task_rewards in per_task_all.items():\n if task_rewards:\n n = len(task_rewards)\n task_mean = sum(task_rewards) / n\n if n > 1:\n task_std = (sum((r - task_mean)**2 for r in task_rewards) / (n - 1)) ** 0.5\n ci_half = 1.96 * task_std / math.sqrt(n)\n else:\n ci_half = 0.0\n log_data[f\"eval/{task_name}\"] = task_mean\n log_data[f\"eval/{task_name}_ci\"] = ci_half\n log_data[f\"eval/{task_name}_n\"] = n\n ci_strs.append(f\"{task_name}={task_mean:.3f}Β±{ci_half:.3f} (n={n})\")\n \n # ββ GDPO per-component stats (Change 5) βββββββββββββββββββββββββββββ\n if per_task_all and all(len(v) > 0 for v in per_task_all.values()):\n try:\n gdpo_rewards = gdpo_normalize(per_task_all)\n log_data[\"eval/gdpo_mean\"] = sum(gdpo_rewards) / len(gdpo_rewards)\n log_data[\"eval/gdpo_std\"] = (sum((r - sum(gdpo_rewards)/len(gdpo_rewards))**2 for r in gdpo_rewards) / len(gdpo_rewards)) ** 0.5\n except Exception as e:\n print(f\" [GDPO] normalization error: {e}\")\n \n # ββ Dynamic task weight update (Change 6) βββββββββββββββββββββββββββ\n per_task_means = {}\n for task_name, task_rewards in per_task_all.items():\n if task_rewards:\n per_task_means[task_name] = sum(task_rewards) / len(task_rewards)\n \n update_task_weights(state.global_step, per_task_means, update_interval=EVAL_STEPS)\n \n for task_name, weight in _task_weights.items():\n log_data[f\"sampler/{task_name}_weight\"] = weight\n \n wandb.log(log_data, step=state.global_step)\n\n status = \"β improved\" if improved else f\"β no gain ({self.no_improve_count + 1}/{self.patience})\"\n print(f\"\\n[EvalReward] step={state.global_step} | mean={mean_reward:.4f} | best={self.best_reward:.4f} | {status}\")\n for cs in ci_strs:\n print(f\" {cs}\")\n print(f\" Task weights: {', '.join(f'{t}={w:.3f}' for t, w in _task_weights.items())}\")\n\n if improved:\n self.best_reward = mean_reward\n self.best_step = state.global_step\n self.no_improve_count = 0\n # ββ V4.2: Save best checkpoint explicitly (Change 8) βββββββββββββ\n best_path = ADAPTER_DIR / \"best_checkpoint\"\n best_path.mkdir(parents=True, exist_ok=True)\n model.save_pretrained(str(best_path))\n tokenizer_local.save_pretrained(str(best_path))\n print(f\" β Best checkpoint saved β {best_path} (reward={mean_reward:.4f})\")\n else:\n self.no_improve_count += 1\n if self.no_improve_count >= self.patience:\n print(f\"[EarlyStopping] No improvement for {self.patience} evals. Halting.\")\n control.should_training_stop = True\n return control\n\n def _run_eval(self, model, tokenizer_local, args):\n FastLanguageModel.for_inference(model)\n rewards = []\n per_task_summary = {\"extraction\": [], \"sql_qa\": [], \"insights\": [], \"push\": []}\n per_task_all = {\"extraction\": [], \"sql_qa\": [], \"insights\": [], \"push\": []}\n \n # V4.2: Use ALL stratified eval samples (65), not just 15\n for record in self.eval_records:\n msgs = record[\"prompt\"]\n text = tokenizer_local.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True)\n user_txt = \" \".join(m.get(\"content\", \"\") for m in msgs if m[\"role\"] == \"user\")\n task = _classify_task_type(user_txt)\n\n inputs = tokenizer_local(text, return_tensors=\"pt\", truncation=True, max_length=args.max_prompt_length).to(model.device)\n with torch.no_grad():\n out = model.generate(\n **inputs,\n max_new_tokens=EVAL_MAX_TOKENS,\n temperature=0.1,\n do_sample=True,\n repetition_penalty=1.0,\n )\n resp = tokenizer_local.decode(out[0][inputs[\"input_ids\"].shape[1]:], skip_special_tokens=True)\n r = self.reward_fn([resp], [text])[0]\n rewards.append(r)\n if task in per_task_all:\n per_task_all[task].append(r)\n per_task_summary[task].append(r)\n \n FastLanguageModel.for_training(model)\n mean_r = sum(rewards) / len(rewards) if rewards else 0.0\n return mean_r, per_task_summary, per_task_all\n\n\n# ββ Training ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ\nFastLanguageModel.for_training(model)\n\ngrpo_config = GRPOConfig(\n output_dir=str(CHECKPOINT_DIR),\n num_generations=NUM_GENERATIONS,\n scale_rewards=SCALE_REWARDS,\n max_completion_length=MAX_COMPLETION_LENGTH,\n max_steps=MAX_STEPS, # V4.2: 1,500\n temperature=TEMPERATURE,\n beta=BETA,\n num_train_epochs=1,\n per_device_train_batch_size=BATCH_SIZE,\n gradient_accumulation_steps=GRAD_ACCUM,\n learning_rate=LEARNING_RATE,\n lr_scheduler_type=LR_SCHEDULER_TYPE,\n warmup_ratio=WARMUP_RATIO,\n fp16=False,\n bf16=True,\n logging_steps=1,\n save_steps=SAVE_STEPS, # V4.2: 100\n save_total_limit=5,\n save_only_model=True,\n report_to=\"wandb\",\n max_prompt_length=MAX_SEQ_LENGTH // 2,\n seed=CURRENT_SEED, # V4.2: per-seed\n remove_unused_columns=False,\n disable_tqdm=True,\n logging_first_step=True,\n)\n\neval_cb = EvalRewardCallbackV2(\n eval_records=list(eval_dataset),\n reward_fn=commerce_reward_fn,\n patience=EARLY_STOPPING_PATIENCE,\n delta=EARLY_STOPPING_DELTA,\n)\n\ntrainer = UnslothGRPOTrainer(\n model=model,\n reward_funcs=commerce_reward_fn,\n args=grpo_config,\n train_dataset=train_dataset,\n processing_class=tokenizer,\n callbacks=[eval_cb],\n)\n\nt_start = time.time()\nresult = trainer.train()\nelapsed = time.time() - t_start\n\nwandb.log({\n \"train/final_loss\": result.training_loss,\n \"train/duration_hours\": elapsed / 3600,\n \"train/total_steps\": result.global_step,\n \"eval/best_reward_final\": eval_cb.best_reward,\n \"eval/best_step\": eval_cb.best_step,\n \"final/task_weights\": _task_weights,\n})\nwandb.finish()\n\nprint(f\"\\n{'='*60}\")\nprint(f\"V4.2 Training Complete (seed={CURRENT_SEED})\")\nprint(f\" Loss: {result.training_loss:.4f}\")\nprint(f\" Steps: {result.global_step}\")\nprint(f\" Duration: {elapsed/3600:.1f}h\")\nprint(f\" Best eval: {eval_cb.best_reward:.4f} (step {eval_cb.best_step})\")\nprint(f\" Final task weights: {_task_weights}\")\nprint(f\"{'='*60}\")"
|
| 163 |
+
},
|
| 164 |
+
{
|
| 165 |
+
"cell_type": "markdown",
|
| 166 |
+
"metadata": {},
|
| 167 |
+
"source": "---\n\n## Cell 14: Post-Training Validation (65 Stratified Samples)\n\n**V4.2:** Full stratified eval with per-task 95% CIs.\n\nReports `mean Β± 1.96 Γ std/βn` for each task.\n\n**The four questions V4.2 must answer:**\n1. Does SQL reward improve with the new reward function?\n2. Is the insights regression noise or forgetting?\n3. Does multi-epoch training push eval above 0.70?\n4. Are results reproducible across seeds?"
|
| 168 |
+
},
|
| 169 |
+
{
|
| 170 |
+
"cell_type": "code",
|
| 171 |
+
"execution_count": null,
|
| 172 |
+
"metadata": {},
|
| 173 |
+
"outputs": [],
|
| 174 |
+
"source": "FastLanguageModel.for_inference(model)\n\nval_samples = list(eval_dataset) # All 65 stratified samples\nval_results = {\"extraction\": [], \"sql_qa\": [], \"insights\": [], \"push\": []}\n\nprint(f\"Post-training validation on {len(val_samples)} stratified samples (seed={CURRENT_SEED})\")\nprint(\"-\" * 80)\n\nfor i, record in enumerate(val_samples):\n msgs = record[\"prompt\"]\n user_text = \" \".join(m[\"content\"] for m in msgs if m[\"role\"] == \"user\")\n task = _classify_task_type(user_text)\n\n text = tokenizer.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True)\n inputs = tokenizer(text, return_tensors=\"pt\").to(model.device)\n with torch.no_grad():\n out = model.generate(\n **inputs,\n max_new_tokens=MAX_COMPLETION_LENGTH,\n temperature=0.1,\n do_sample=True,\n repetition_penalty=1.0,\n )\n resp = tokenizer.decode(out[0][inputs[\"input_ids\"].shape[1]:], skip_special_tokens=True)\n r = commerce_reward_fn([resp], [text])[0]\n val_results[task].append(r)\n if i < 10 or r < 0.2: # Print first 10 and all low-scoring\n print(f\" [{task:12s}] reward={r:.3f} | {strip_think(resp)[:80]}\")\n\n# ββ Results with 95% CIs ββββββββββββββββββββββββββββββββββββββββββββββββββββ\nprint(f\"\\n{'='*80}\")\nprint(f\"VALIDATION RESULTS β V4.2 Seed {CURRENT_SEED}\")\nprint(f\"{'='*80}\")\nprint(f\"{'Task':15s} {'Mean':>8s} {'Β± 95% CI':>10s} {'Min':>6s} {'Max':>6s} {'N':>4s}\")\nprint(\"-\" * 55)\n\noverall = []\nresults_by_seed = {} # Store for cross-seed comparison\n\nfor task in [\"extraction\", \"sql_qa\", \"insights\", \"push\"]:\n rewards = val_results[task]\n overall.extend(rewards)\n if rewards:\n n = len(rewards)\n mean_r = sum(rewards) / n\n if n > 1:\n std_r = (sum((r - mean_r)**2 for r in rewards) / (n - 1)) ** 0.5\n ci_half = 1.96 * std_r / math.sqrt(n)\n else:\n std_r = 0.0\n ci_half = 0.0\n print(f\"{task:15s} {mean_r:8.3f} {'Β±':>2s}{ci_half:7.3f} {min(rewards):6.3f} {max(rewards):6.3f} {n:4d}\")\n results_by_seed[task] = {\"mean\": mean_r, \"ci\": ci_half, \"n\": n, \"std\": std_r}\n\nif overall:\n n_total = len(overall)\n mean_total = sum(overall) / n_total\n std_total = (sum((r - mean_total)**2 for r in overall) / (n_total - 1)) ** 0.5\n ci_total = 1.96 * std_total / math.sqrt(n_total)\n print(\"-\" * 55)\n print(f\"{'OVERALL':15s} {mean_total:8.3f} {'Β±':>2s}{ci_total:7.3f} {min(overall):6.3f} {max(overall):6.3f} {n_total:4d}\")\n results_by_seed[\"overall\"] = {\"mean\": mean_total, \"ci\": ci_total, \"n\": n_total, \"std\": std_total}\n\n# ββ Save results for cross-seed comparison ββββββββββββββββββββββββββββββββββ\nresults_file = ADAPTER_DIR / f\"eval_results_seed{CURRENT_SEED}.json\"\nresults_file.parent.mkdir(parents=True, exist_ok=True)\nwith open(results_file, \"w\") as f:\n json.dump(results_by_seed, f, indent=2)\nprint(f\"\\nβ Results saved to {results_file}\")\n\n# ββ V4.2 Decision βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ\nprint(f\"\\n--- V4.2 Questions ---\")\nsql_mean = results_by_seed.get(\"sql_qa\", {}).get(\"mean\", 0)\ninsights_mean = results_by_seed.get(\"insights\", {}).get(\"mean\", 0)\noverall_mean = results_by_seed.get(\"overall\", {}).get(\"mean\", 0)\n\nprint(f\"Q1 SQL reward: {sql_mean:.3f} ({'improved' if sql_mean > 0.60 else 'still stagnant' if sql_mean < 0.56 else 'modest gain'})\")\nprint(f\"Q2 Insights: {insights_mean:.3f} ({'stable' if insights_mean > 0.70 else 'regressed' if insights_mean < 0.60 else 'mixed'})\")\nprint(f\"Q3 Overall: {overall_mean:.3f} ({'above 0.70 target' if overall_mean > 0.70 else 'below target'})\")\nprint(f\"Q4 Seeds: Seed {CURRENT_SEED} done. Run seeds {[s for s in SEEDS if s != CURRENT_SEED]} next.\")"
|
| 175 |
+
},
|
| 176 |
+
{
|
| 177 |
+
"cell_type": "markdown",
|
| 178 |
+
"metadata": {},
|
| 179 |
+
"source": "---\n\n## Cell 15: Save Adapter\n\n**V4.2:** Saves from `best_checkpoint/` (peak eval reward) instead of last training step."
|
| 180 |
+
},
|
| 181 |
+
{
|
| 182 |
+
"cell_type": "code",
|
| 183 |
+
"execution_count": null,
|
| 184 |
+
"metadata": {},
|
| 185 |
+
"outputs": [],
|
| 186 |
+
"source": "# V4.2: Save the best checkpoint, not the last training step\nbest_checkpoint_path = ADAPTER_DIR / \"best_checkpoint\"\n\nif best_checkpoint_path.exists():\n print(f\"β Best checkpoint exists at {best_checkpoint_path}\")\n print(f\" Best eval reward: {eval_cb.best_reward:.4f} (step {eval_cb.best_step})\")\n # Copy best checkpoint to main adapter dir for easy loading\n import shutil\n final_path = ADAPTER_DIR / \"final\"\n if final_path.exists():\n shutil.rmtree(final_path)\n shutil.copytree(str(best_checkpoint_path), str(final_path))\n print(f\" β Copied to {final_path}\")\nelse:\n print(\"β οΈ No best_checkpoint found. Saving current model state.\")\n ADAPTER_DIR.mkdir(parents=True, exist_ok=True)\n model.save_pretrained(str(ADAPTER_DIR / \"final\"))\n tokenizer.save_pretrained(str(ADAPTER_DIR / \"final\"))\n\nprint(f\"\\nβ Adapter saved for seed {CURRENT_SEED}\")\nprint(f\" Location: {ADAPTER_DIR / 'final'}\")\nprint(f\" Best eval: {eval_cb.best_reward:.4f} at step {eval_cb.best_step}\")"
|
| 187 |
+
},
|
| 188 |
+
{
|
| 189 |
+
"cell_type": "markdown",
|
| 190 |
+
"metadata": {},
|
| 191 |
+
"source": "---\n\n## Cell 16: Results Table Generation (Change 7)\n\n**Run this after ALL 3 seeds are complete.**\n\nReads `eval_results_seedN.json` from each seed directory and produces the cross-seed\nresults table with mean Β± 95% CI.\n\n```\n| Task | Seed 42 | Seed 123 | Seed 456 | Mean Β± 95% CI |\n|---|---|---|---|---|\n| Extraction | ... | ... | ... | X.XX Β± 0.0X |\n| SQL Q&A | ... | ... | ... | X.XX Β± 0.0X |\n| Insights | ... | ... | ... | X.XX Β± 0.0X |\n| Push | ... | ... | ... | X.XX Β± 0.0X |\n| **Mean** | ... | ... | ... | **X.XX Β± 0.0X** |\n```"
|
| 192 |
+
},
|
| 193 |
+
{
|
| 194 |
+
"cell_type": "code",
|
| 195 |
+
"execution_count": null,
|
| 196 |
+
"metadata": {},
|
| 197 |
+
"outputs": [],
|
| 198 |
+
"source": "# ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ\n# V4.2: Cross-Seed Results Table (Change 7)\n# Run this AFTER completing all 3 seeds (42, 123, 456)\n# ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ\n\nresults_by_seed = {}\nmissing_seeds = []\n\nfor seed in SEEDS:\n seed_dir = MODELS_DIR / f\"tucano2-0.5B-instruct-grpo-v4.2-seed{seed}\"\n results_file = seed_dir / f\"eval_results_seed{seed}.json\"\n if results_file.exists():\n with open(results_file) as f:\n results_by_seed[seed] = json.load(f)\n print(f\"β Loaded results for seed {seed}\")\n else:\n missing_seeds.append(seed)\n print(f\"β οΈ Missing results for seed {seed} (run the notebook with CURRENT_SEED={seed})\")\n\nif missing_seeds:\n print(f\"\\nβ οΈ Missing seeds: {missing_seeds}\")\n print(f\" Complete these runs before generating the final table.\")\n print(f\" Change CURRENT_SEED in Cell 3 and re-run Cells 3-15.\")\n\nif len(results_by_seed) >= 2:\n print(f\"\\n{'='*90}\")\n print(f\"V4.2 CROSS-SEED RESULTS TABLE\")\n print(f\"{'='*90}\")\n \n tasks = [\"extraction\", \"sql_qa\", \"insights\", \"push\", \"overall\"]\n \n # Header\n header = f\"{'Task':15s}\"\n for seed in SEEDS:\n if seed in results_by_seed:\n header += f\" {'Seed '+str(seed):>10s}\"\n header += f\" {'Mean Β± 95% CI':>18s}\"\n print(header)\n print(\"-\" * len(header))\n \n for task in tasks:\n row = f\"{task:15s}\"\n seed_means = []\n for seed in SEEDS:\n if seed in results_by_seed and task in results_by_seed[seed]:\n m = results_by_seed[seed][task][\"mean\"]\n seed_means.append(m)\n row += f\" {m:10.3f}\"\n elif seed in results_by_seed:\n row += f\" {'β':>10s}\"\n \n if len(seed_means) >= 2:\n cross_mean = sum(seed_means) / len(seed_means)\n cross_std = (sum((m - cross_mean)**2 for m in seed_means) / (len(seed_means) - 1)) ** 0.5\n # With 3 seeds, use t-distribution critical value (t=4.303 for 95% CI, df=2)\n # But for consistency with the handoff, use Β±std\n row += f\" {cross_mean:7.3f} Β± {cross_std:.3f}\"\n elif len(seed_means) == 1:\n row += f\" {seed_means[0]:7.3f} (1 seed)\"\n \n if task == \"overall\":\n row = f\"**{row.strip()}**\"\n print(row)\n \n print(f\"\\n{'='*90}\")\n \n # ββ Reproducibility assessment ββββββββββββββββββββββββββββββββββββββββββ\n if len(results_by_seed) == 3:\n overall_means = [results_by_seed[s][\"overall\"][\"mean\"] for s in SEEDS if s in results_by_seed]\n overall_std = (sum((m - sum(overall_means)/len(overall_means))**2 for m in overall_means) / (len(overall_means) - 1)) ** 0.5\n print(f\"\\nReproducibility: overall std = {overall_std:.4f}\")\n if overall_std < 0.03:\n print(f\" β
Robust (std < 0.03): results are reproducible across seeds\")\n elif overall_std < 0.05:\n print(f\" β οΈ Moderate (0.03 < std < 0.05): some initialization sensitivity\")\n else:\n print(f\" β High variance (std > 0.05): significant initialization sensitivity\")\n\nelse:\n print(f\"\\nNeed at least 2 seeds to generate comparison table.\")\n print(f\"Current seed ({CURRENT_SEED}) results:\")\n if CURRENT_SEED in results_by_seed:\n for task, data in results_by_seed[CURRENT_SEED].items():\n print(f\" {task}: {data['mean']:.3f} Β± {data.get('ci', 0):.3f}\")"
|
| 199 |
+
}
|
| 200 |
+
],
|
| 201 |
+
"metadata": {
|
| 202 |
+
"kernelspec": {
|
| 203 |
+
"display_name": "Python 3",
|
| 204 |
+
"language": "python",
|
| 205 |
+
"name": "python3"
|
| 206 |
+
},
|
| 207 |
+
"language_info": {
|
| 208 |
+
"name": "python",
|
| 209 |
+
"version": "3.10.0"
|
| 210 |
+
}
|
| 211 |
+
},
|
| 212 |
+
"nbformat": 4,
|
| 213 |
+
"nbformat_minor": 5
|
| 214 |
+
}
|