fix(notebook): correct reward weights, fix compose_reward_single signature + tuple unpack, add model hint

notebooks/fathom_train.ipynb

@@ -87,16 +87,39 @@
    "source": [
     "from rewards.compose import compose_reward_single\n",
     "import types\n",
-    "cfg = types.SimpleNamespace(alpha=0.2, weights=types.SimpleNamespace(correctness=0.75, token_budget=0.2, recursion_efficiency=0.05), token_budget_variant='capped_linear', answer_regex='<answer>(.*?)</answer>')\n",
     "\n",
-    "good = '<answer>Paris</answer>'\n",
+    "# Weights match configs/reward/v1.yaml exactly (REW-02 v3)\n",
+    "cfg = types.SimpleNamespace(\n",
+    "    alpha=0.2,\n",
+    "    weights=types.SimpleNamespace(\n",
+    "        correctness=0.70,\n",
+    "        token_budget=0.15,\n",
+    "        recursion_efficiency=0.15\n",
+    "    ),\n",
+    "    token_budget_variant='capped_linear',\n",
+    "    answer_regex='<answer>(.*?)</answer>'\n",
+    ")\n",
+    "\n",
+    "good             = '<answer>Paris</answer>'\n",
     "format_only_wrong = '<answer>Berlin</answer>'\n",
-    "no_format = 'Paris'\n",
-    "padded = '<answer>Paris</answer>' + ' '*5000\n",
+    "no_format        = 'Paris'\n",
+    "padded           = '<answer>Paris</answer>' + ' '*5000\n",
     "\n",
-    "for label, completion in [('correct', good), ('wrong-but-formatted', format_only_wrong), ('no-format', no_format), ('length-padded', padded)]:\n",
-    "    score = compose_reward_single(completion=completion, gold_answer='Paris', prompt_token_count=512, llm_call_count=0, cfg_reward=cfg)\n",
-    "    print(f'{label:25s} -> reward = {score:.4f}')"
+    "for label, completion in [\n",
+    "    ('correct',             good),\n",
+    "    ('wrong-but-formatted', format_only_wrong),\n",
+    "    ('no-format',           no_format),\n",
+    "    ('length-padded',       padded),\n",
+    "]:\n",
+    "    # compose_reward_single returns (composite_score, metrics_dict)\n",
+    "    score, metrics = compose_reward_single(\n",
+    "        completion=completion,\n",
+    "        gold_answer='Paris',\n",
+    "        prompt_token_count=512,\n",
+    "        cfg_reward=cfg,\n",
+    "        llm_call_count=0,\n",
+    "    )\n",
+    "    print(f'{label:25s} -> reward = {score:.4f}  (correctness={metrics[\"correctness\"]:.2f})')"
    ]
   },
   {
@@ -136,7 +159,10 @@
     "# with initialize(config_path='configs', version_base='1.3'):\n",
     "#     cfg = compose(config_name='config', overrides=['model=qwen_1_5b','train=grpo','train.max_steps=50'])\n",
     "# m, t = load_model_and_tokenizer(cfg)\n",
-    "# run_grpo(cfg, m, t, make_reward_fn(cfg.reward), 'https://Pratham-math-fathom-env.hf.space')"
+    "# run_grpo(cfg, m, t, make_reward_fn(cfg.reward), 'https://Pratham-math-fathom-env.hf.space')\n",
+    "#\n",
+    "# NOTE: Default config uses qwen_1_5b (1.5B Qwen2.5-Coder-Instruct-bnb-4bit).\n",
+    "# For a faster smoke run on smaller GPU, override with model=qwen_0_5b_smoke."
    ]
   },
   {