heavycoderhh commited on
Commit
a6670df
·
verified ·
1 Parent(s): 83c8e30

Add run4b QLoRA SFT training notebook

Browse files
Files changed (1) hide show
  1. notebooks/train_counsel_run4b.ipynb +514 -0
notebooks/train_counsel_run4b.ipynb ADDED
@@ -0,0 +1,514 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "cells": [
3
+ {
4
+ "cell_type": "markdown",
5
+ "metadata": {},
6
+ "source": [
7
+ "# Counsel-Env Run4b: 4-bit QLoRA Oracle-SFT on Qwen3-8B\n",
8
+ "\n",
9
+ "This is the official training notebook for the **run4b** checkpoint:\n",
10
+ "\n",
11
+ "- HF training job ID: `69edb014d2c8bd8662bcf5ba`\n",
12
+ "- Base model: `Qwen/Qwen3-8B`\n",
13
+ "- Method: 4-bit QLoRA SFT on assistant-only oracle next-action rows\n",
14
+ "- Output adapter: [`heavycoderhh/counsel-env-qwen3-8b-qlora-sft-run4b`](https://huggingface.co/heavycoderhh/counsel-env-qwen3-8b-qlora-sft-run4b)\n",
15
+ "- Eval result: `avg_reward=0.864`, `surface_rate=0.943` on the 150-seed held-out eval\n",
16
+ "\n",
17
+ "It mirrors `counsel_env/scripts/run_qlora_sft_training_job.py`, the exact script used to launch the HF Job that produced run4b. Run it locally on a single A100 / H100, on a Colab Pro A100 instance, or via `hf jobs run`. Each oracle row teaches the model the cross-examination pattern: trigger a sealed witness claim, then present the matching disprover exhibit."
18
+ ]
19
+ },
20
+ {
21
+ "cell_type": "markdown",
22
+ "metadata": {},
23
+ "source": [
24
+ "## 1. Install dependencies\n",
25
+ "\n",
26
+ "Pinned to the same versions used inside the HF Job that produced run4b. Skip this cell if you already have a compatible environment."
27
+ ]
28
+ },
29
+ {
30
+ "cell_type": "code",
31
+ "execution_count": null,
32
+ "metadata": {},
33
+ "outputs": [],
34
+ "source": [
35
+ "# !pip install -q \\\n",
36
+ "# 'accelerate>=1.12.0' \\\n",
37
+ "# 'bitsandbytes>=0.48.0' \\\n",
38
+ "# 'datasets>=4.0.0' \\\n",
39
+ "# 'huggingface_hub>=1.0.0' \\\n",
40
+ "# 'openenv-core>=0.2.1' \\\n",
41
+ "# 'peft>=0.18.0' \\\n",
42
+ "# 'torch>=2.8.0' \\\n",
43
+ "# 'transformers>=5.2.0'"
44
+ ]
45
+ },
46
+ {
47
+ "cell_type": "markdown",
48
+ "metadata": {},
49
+ "source": [
50
+ "## 2. Configuration\n",
51
+ "\n",
52
+ "Defaults reproduce run4b exactly. Override any value with environment variables before launching the cell, e.g. `os.environ['COUNSEL_SFT_MAX_STEPS'] = '10'` for a smoke test."
53
+ ]
54
+ },
55
+ {
56
+ "cell_type": "code",
57
+ "execution_count": null,
58
+ "metadata": {},
59
+ "outputs": [],
60
+ "source": [
61
+ "from __future__ import annotations\n",
62
+ "\n",
63
+ "import importlib.util\n",
64
+ "import json\n",
65
+ "import os\n",
66
+ "import sys\n",
67
+ "from pathlib import Path\n",
68
+ "from typing import Any, Dict, List\n",
69
+ "\n",
70
+ "import torch\n",
71
+ "from datasets import Dataset\n",
72
+ "from huggingface_hub import HfApi, snapshot_download\n",
73
+ "from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training\n",
74
+ "from transformers import (\n",
75
+ " AutoModelForCausalLM,\n",
76
+ " AutoTokenizer,\n",
77
+ " BitsAndBytesConfig,\n",
78
+ " DataCollatorForSeq2Seq,\n",
79
+ " Trainer,\n",
80
+ " TrainingArguments,\n",
81
+ ")\n",
82
+ "\n",
83
+ "SPACE_REPO = os.getenv('COUNSEL_SPACE_REPO', 'heavycoderhh/counsel-env')\n",
84
+ "MODEL = os.getenv('COUNSEL_MODEL', 'Qwen/Qwen3-8B')\n",
85
+ "OUTPUT_DIR = Path(os.getenv('COUNSEL_OUTPUT_DIR', '/tmp/counsel-qwen3-8b-qlora-sft'))\n",
86
+ "ARTIFACT_REPO = os.getenv('COUNSEL_ARTIFACT_REPO', 'heavycoderhh/counsel-env-qwen3-8b-qlora-sft-run4b')\n",
87
+ "SFT_DATASET_SIZE = int(os.getenv('COUNSEL_SFT_DATASET_SIZE', '480'))\n",
88
+ "SFT_EPOCHS = float(os.getenv('COUNSEL_SFT_EPOCHS', '1'))\n",
89
+ "SFT_MAX_STEPS = int(os.getenv('COUNSEL_SFT_MAX_STEPS', '220'))\n",
90
+ "SFT_LEARNING_RATE = float(os.getenv('COUNSEL_SFT_LEARNING_RATE', '1e-4'))\n",
91
+ "MAX_SFT_LENGTH = int(os.getenv('COUNSEL_MAX_SFT_LENGTH', '1536'))\n",
92
+ "LORA_R = int(os.getenv('COUNSEL_LORA_R', '16'))\n",
93
+ "LORA_ALPHA = int(os.getenv('COUNSEL_LORA_ALPHA', '32'))\n",
94
+ "LORA_DROPOUT = float(os.getenv('COUNSEL_LORA_DROPOUT', '0.05'))\n",
95
+ "GRADIENT_ACCUMULATION_STEPS = int(os.getenv('COUNSEL_GRAD_ACCUM', '4'))\n",
96
+ "INCLUDE_REST_ROWS = os.getenv('COUNSEL_INCLUDE_REST_ROWS', '0') == '1'\n",
97
+ "\n",
98
+ "print({\n",
99
+ " 'model': MODEL,\n",
100
+ " 'artifact_repo': ARTIFACT_REPO,\n",
101
+ " 'sft_dataset_size': SFT_DATASET_SIZE,\n",
102
+ " 'sft_max_steps': SFT_MAX_STEPS,\n",
103
+ " 'sft_learning_rate': SFT_LEARNING_RATE,\n",
104
+ " 'lora_r': LORA_R,\n",
105
+ " 'lora_alpha': LORA_ALPHA,\n",
106
+ " 'gradient_accumulation_steps': GRADIENT_ACCUMULATION_STEPS,\n",
107
+ "})"
108
+ ]
109
+ },
110
+ {
111
+ "cell_type": "markdown",
112
+ "metadata": {},
113
+ "source": [
114
+ "## 3. Load the Counsel-Env package\n",
115
+ "\n",
116
+ "Pulls the latest Counsel-Env source from the HF Space when the package is not already importable (matches how the HF Job ran)."
117
+ ]
118
+ },
119
+ {
120
+ "cell_type": "code",
121
+ "execution_count": null,
122
+ "metadata": {},
123
+ "outputs": [],
124
+ "source": [
125
+ "def prepare_imports() -> None:\n",
126
+ " try:\n",
127
+ " import counsel_env # noqa: F401\n",
128
+ " return\n",
129
+ " except Exception:\n",
130
+ " pass\n",
131
+ " source_dir = snapshot_download(repo_id=SPACE_REPO, repo_type='space')\n",
132
+ " init_path = Path(source_dir) / '__init__.py'\n",
133
+ " spec = importlib.util.spec_from_file_location(\n",
134
+ " 'counsel_env',\n",
135
+ " init_path,\n",
136
+ " submodule_search_locations=[source_dir],\n",
137
+ " )\n",
138
+ " if spec is None or spec.loader is None:\n",
139
+ " raise RuntimeError(f'Could not load Counsel-Env package from {source_dir}')\n",
140
+ " module = importlib.util.module_from_spec(spec)\n",
141
+ " sys.modules['counsel_env'] = module\n",
142
+ " spec.loader.exec_module(module)\n",
143
+ " print(f'Loaded Counsel-Env source from {source_dir}')\n",
144
+ "\n",
145
+ "prepare_imports()\n",
146
+ "\n",
147
+ "from counsel_env.models import CounselAction\n",
148
+ "from counsel_env.server.counsel_env_environment import CounselEnvironment"
149
+ ]
150
+ },
151
+ {
152
+ "cell_type": "markdown",
153
+ "metadata": {},
154
+ "source": [
155
+ "## 4. Tool schema and prompt\n",
156
+ "\n",
157
+ "These functions are used only as Qwen tool schemas (the model emits structured `<tool_call>` blocks). They never execute at training time."
158
+ ]
159
+ },
160
+ {
161
+ "cell_type": "code",
162
+ "execution_count": null,
163
+ "metadata": {},
164
+ "outputs": [],
165
+ "source": [
166
+ "def ask_question(question: str) -> str:\n",
167
+ " \"\"\"Ask the witness a question.\n",
168
+ "\n",
169
+ " Args:\n",
170
+ " question: The cross-examination question to ask the witness.\n",
171
+ " \"\"\"\n",
172
+ " raise RuntimeError('Tool schema only')\n",
173
+ "\n",
174
+ "\n",
175
+ "def present_evidence(exhibit_id: str) -> str:\n",
176
+ " \"\"\"Present an exhibit to the witness.\n",
177
+ "\n",
178
+ " Args:\n",
179
+ " exhibit_id: The ID of the exhibit to present.\n",
180
+ " \"\"\"\n",
181
+ " raise RuntimeError('Tool schema only')\n",
182
+ "\n",
183
+ "\n",
184
+ "def rest_case() -> str:\n",
185
+ " \"\"\"End the cross-examination.\"\"\"\n",
186
+ " raise RuntimeError('Tool schema only')\n",
187
+ "\n",
188
+ "\n",
189
+ "TOOLS = [ask_question, present_evidence, rest_case]\n",
190
+ "\n",
191
+ "BASE_PROMPT = (\n",
192
+ " 'You are a sharp prosecuting attorney cross-examining a deterministic witness. '\n",
193
+ " 'Your goal is to surface contradictions by first making the witness commit to a claim, '\n",
194
+ " 'then presenting the exact exhibit that disproves it. Use the limited question budget efficiently. '\n",
195
+ " 'Return exactly one tool call and no prose. Never invent exhibit IDs. '\n",
196
+ " 'Do not rest the case until after at least one contradiction has been surfaced.'\n",
197
+ ")"
198
+ ]
199
+ },
200
+ {
201
+ "cell_type": "markdown",
202
+ "metadata": {},
203
+ "source": [
204
+ "## 5. Curriculum schedule and rendering helpers\n",
205
+ "\n",
206
+ "Same 45/40/15 easy-medium-hard split used in the run4b job. Each generated case becomes one or more `(prompt, assistant)` SFT rows."
207
+ ]
208
+ },
209
+ {
210
+ "cell_type": "code",
211
+ "execution_count": null,
212
+ "metadata": {},
213
+ "outputs": [],
214
+ "source": [
215
+ "def sample_stage_schedule(total: int) -> List[str]:\n",
216
+ " schedule = []\n",
217
+ " for index in range(total):\n",
218
+ " frac = index / max(1, total - 1)\n",
219
+ " if frac < 0.45:\n",
220
+ " schedule.append('easy')\n",
221
+ " elif frac < 0.85:\n",
222
+ " schedule.append('medium')\n",
223
+ " else:\n",
224
+ " schedule.append('hard')\n",
225
+ " return schedule\n",
226
+ "\n",
227
+ "\n",
228
+ "def format_evidence(evidence_descriptions: Dict[str, str]) -> str:\n",
229
+ " return '\\n'.join(f'- {exhibit_id}: {description}' for exhibit_id, description in evidence_descriptions.items())\n",
230
+ "\n",
231
+ "\n",
232
+ "def reset_text(obs: Any) -> str:\n",
233
+ " return (\n",
234
+ " f'CASE BRIEF:\\n{obs.case_brief}\\n\\n'\n",
235
+ " f'You have {obs.questions_remaining} questions. '\n",
236
+ " 'Available exhibits with descriptions:\\n'\n",
237
+ " f'{format_evidence(obs.evidence_descriptions)}\\n\\n'\n",
238
+ " 'Use ask_question first. After the witness commits, present the exact matching exhibit ID.'\n",
239
+ " )\n",
240
+ "\n",
241
+ "\n",
242
+ "def tool_feedback(obs: Any) -> str:\n",
243
+ " components = obs.reward_components or {}\n",
244
+ " evidence = ', '.join(obs.available_evidence)\n",
245
+ " return (\n",
246
+ " f'WITNESS: {obs.witness_response}\\n'\n",
247
+ " f\"STATE: triggered={int(components.get('contradictions_triggered', 0))}/\"\n",
248
+ " f\"{int(components.get('contradictions_total', 0))}, \"\n",
249
+ " f\"surfaced={int(components.get('contradictions_surfaced', 0))}/\"\n",
250
+ " f\"{int(components.get('contradictions_total', 0))}, \"\n",
251
+ " f'questions_remaining={obs.questions_remaining}, done={obs.done}\\n'\n",
252
+ " f'VALID_EXHIBITS: {evidence}'\n",
253
+ " )\n",
254
+ "\n",
255
+ "\n",
256
+ "def tool_call(name: str, arguments: Dict[str, Any]) -> str:\n",
257
+ " return '<tool_call>\\n' + json.dumps({'name': name, 'arguments': arguments}, sort_keys=True) + '\\n</tool_call>'"
258
+ ]
259
+ },
260
+ {
261
+ "cell_type": "markdown",
262
+ "metadata": {},
263
+ "source": [
264
+ "## 6. Build the assistant-only oracle SFT dataset\n",
265
+ "\n",
266
+ "For each generated case we walk the deterministic environment with the oracle policy: trigger keyword question, then matching disprover exhibit. Loss is masked to the assistant span only, so the model is graded purely on producing the correct tool call."
267
+ ]
268
+ },
269
+ {
270
+ "cell_type": "code",
271
+ "execution_count": null,
272
+ "metadata": {},
273
+ "outputs": [],
274
+ "source": [
275
+ "def tokenize_assistant_action(\n",
276
+ " tokenizer: Any,\n",
277
+ " prompt_messages: List[Dict[str, str]],\n",
278
+ " assistant_message: Dict[str, str],\n",
279
+ ") -> Dict[str, List[int]]:\n",
280
+ " prompt_text = tokenizer.apply_chat_template(\n",
281
+ " prompt_messages,\n",
282
+ " tools=TOOLS,\n",
283
+ " tokenize=False,\n",
284
+ " add_generation_prompt=True,\n",
285
+ " chat_template_kwargs={'enable_thinking': False},\n",
286
+ " )\n",
287
+ " full_text = tokenizer.apply_chat_template(\n",
288
+ " prompt_messages + [assistant_message],\n",
289
+ " tools=TOOLS,\n",
290
+ " tokenize=False,\n",
291
+ " add_generation_prompt=False,\n",
292
+ " chat_template_kwargs={'enable_thinking': False},\n",
293
+ " )\n",
294
+ " encoded = tokenizer(full_text, add_special_tokens=False, max_length=MAX_SFT_LENGTH, truncation=True)\n",
295
+ " prompt_ids = tokenizer(prompt_text, add_special_tokens=False, max_length=MAX_SFT_LENGTH, truncation=True)['input_ids']\n",
296
+ " input_ids = encoded['input_ids']\n",
297
+ " labels = list(input_ids)\n",
298
+ " labels[: min(len(prompt_ids), len(labels))] = [-100] * min(len(prompt_ids), len(labels))\n",
299
+ " encoded['labels'] = labels\n",
300
+ " return encoded\n",
301
+ "\n",
302
+ "\n",
303
+ "def create_oracle_sft_dataset(tokenizer: Any, num_samples: int) -> Dataset:\n",
304
+ " rows: List[Dict[str, List[int]]] = []\n",
305
+ " stages = sample_stage_schedule(num_samples)\n",
306
+ " for seed, stage in enumerate(stages, start=514159):\n",
307
+ " env = CounselEnvironment()\n",
308
+ " obs = env.reset(seed=seed, curriculum_stage=stage, episode_id=f'qwen8b_sft_{seed}')\n",
309
+ " messages = [{'role': 'user', 'content': f'{BASE_PROMPT}\\n\\n{reset_text(obs)}'}]\n",
310
+ " for contradiction in env.witness.contradictions:\n",
311
+ " question = f'{contradiction.trigger_keywords[0]}?'\n",
312
+ " assistant_message = {'role': 'assistant', 'content': tool_call('ask_question', {'question': question})}\n",
313
+ " rows.append(tokenize_assistant_action(tokenizer, messages, assistant_message))\n",
314
+ " messages.append(assistant_message)\n",
315
+ " obs = env.step(CounselAction(tool='ask_question', text=question))\n",
316
+ " messages.append({'role': 'user', 'content': f'<tool_response>\\n{tool_feedback(obs)}\\n</tool_response>'})\n",
317
+ " exhibit_id = contradiction.disprover_evidence_id\n",
318
+ " assistant_message = {\n",
319
+ " 'role': 'assistant',\n",
320
+ " 'content': tool_call('present_evidence', {'exhibit_id': exhibit_id}),\n",
321
+ " }\n",
322
+ " rows.append(tokenize_assistant_action(tokenizer, messages, assistant_message))\n",
323
+ " messages.append(assistant_message)\n",
324
+ " obs = env.step(CounselAction(tool='present_evidence', exhibit_id=exhibit_id))\n",
325
+ " messages.append({'role': 'user', 'content': f'<tool_response>\\n{tool_feedback(obs)}\\n</tool_response>'})\n",
326
+ " if INCLUDE_REST_ROWS:\n",
327
+ " assistant_message = {'role': 'assistant', 'content': tool_call('rest_case', {})}\n",
328
+ " rows.append(tokenize_assistant_action(tokenizer, messages, assistant_message))\n",
329
+ " return Dataset.from_list(rows)"
330
+ ]
331
+ },
332
+ {
333
+ "cell_type": "markdown",
334
+ "metadata": {},
335
+ "source": [
336
+ "## 7. Load the 4-bit model and attach LoRA adapter\n",
337
+ "\n",
338
+ "NF4 quantization with bf16 compute, double quantization, and `paged_adamw_8bit` optimizer. LoRA targets every Qwen attention and MLP projection."
339
+ ]
340
+ },
341
+ {
342
+ "cell_type": "code",
343
+ "execution_count": null,
344
+ "metadata": {},
345
+ "outputs": [],
346
+ "source": [
347
+ "tokenizer = AutoTokenizer.from_pretrained(MODEL, trust_remote_code=True)\n",
348
+ "if tokenizer.pad_token is None:\n",
349
+ " tokenizer.pad_token = tokenizer.eos_token\n",
350
+ "\n",
351
+ "dataset = create_oracle_sft_dataset(tokenizer, SFT_DATASET_SIZE)\n",
352
+ "print(f'sft_row_count={len(dataset)}')\n",
353
+ "\n",
354
+ "quantization = BitsAndBytesConfig(\n",
355
+ " load_in_4bit=True,\n",
356
+ " bnb_4bit_quant_type='nf4',\n",
357
+ " bnb_4bit_compute_dtype=torch.bfloat16,\n",
358
+ " bnb_4bit_use_double_quant=True,\n",
359
+ ")\n",
360
+ "model = AutoModelForCausalLM.from_pretrained(\n",
361
+ " MODEL,\n",
362
+ " quantization_config=quantization,\n",
363
+ " device_map='auto',\n",
364
+ " trust_remote_code=True,\n",
365
+ ")\n",
366
+ "model.config.use_cache = False\n",
367
+ "model = prepare_model_for_kbit_training(model)\n",
368
+ "lora_config = LoraConfig(\n",
369
+ " r=LORA_R,\n",
370
+ " lora_alpha=LORA_ALPHA,\n",
371
+ " lora_dropout=LORA_DROPOUT,\n",
372
+ " bias='none',\n",
373
+ " task_type='CAUSAL_LM',\n",
374
+ " target_modules=['q_proj', 'k_proj', 'v_proj', 'o_proj', 'gate_proj', 'up_proj', 'down_proj'],\n",
375
+ ")\n",
376
+ "model = get_peft_model(model, lora_config)\n",
377
+ "model.print_trainable_parameters()"
378
+ ]
379
+ },
380
+ {
381
+ "cell_type": "markdown",
382
+ "metadata": {},
383
+ "source": [
384
+ "## 8. Train\n",
385
+ "\n",
386
+ "Default budget: 220 steps, batch size 1 with grad-accum 4 (effective batch size 4), learning rate 1e-4. Run4b finished this in **1287.7 seconds** with `train_loss = 0.0565`."
387
+ ]
388
+ },
389
+ {
390
+ "cell_type": "code",
391
+ "execution_count": null,
392
+ "metadata": {},
393
+ "outputs": [],
394
+ "source": [
395
+ "OUTPUT_DIR.mkdir(parents=True, exist_ok=True)\n",
396
+ "\n",
397
+ "args = TrainingArguments(\n",
398
+ " output_dir=str(OUTPUT_DIR),\n",
399
+ " per_device_train_batch_size=1,\n",
400
+ " gradient_accumulation_steps=GRADIENT_ACCUMULATION_STEPS,\n",
401
+ " num_train_epochs=SFT_EPOCHS,\n",
402
+ " max_steps=SFT_MAX_STEPS if SFT_MAX_STEPS > 0 else -1,\n",
403
+ " learning_rate=SFT_LEARNING_RATE,\n",
404
+ " optim='paged_adamw_8bit',\n",
405
+ " logging_steps=10,\n",
406
+ " save_strategy='no',\n",
407
+ " report_to='none',\n",
408
+ " bf16=torch.cuda.is_available(),\n",
409
+ " gradient_checkpointing=True,\n",
410
+ " remove_unused_columns=False,\n",
411
+ ")\n",
412
+ "trainer = Trainer(\n",
413
+ " model=model,\n",
414
+ " args=args,\n",
415
+ " train_dataset=dataset,\n",
416
+ " data_collator=DataCollatorForSeq2Seq(tokenizer=tokenizer, padding=True, label_pad_token_id=-100),\n",
417
+ ")\n",
418
+ "train_result = trainer.train()\n",
419
+ "trainer.save_model(str(OUTPUT_DIR))\n",
420
+ "tokenizer.save_pretrained(str(OUTPUT_DIR))\n",
421
+ "print(train_result.metrics)"
422
+ ]
423
+ },
424
+ {
425
+ "cell_type": "markdown",
426
+ "metadata": {},
427
+ "source": [
428
+ "## 9. Write training summary and (optionally) upload the adapter\n",
429
+ "\n",
430
+ "The summary JSON is exactly the file mirrored at `assets/trained_eval_run4b_8b_sft/training_summary.json` for run4b. Adapter upload runs only when `HF_TOKEN` is set."
431
+ ]
432
+ },
433
+ {
434
+ "cell_type": "code",
435
+ "execution_count": null,
436
+ "metadata": {},
437
+ "outputs": [],
438
+ "source": [
439
+ "summary = {\n",
440
+ " 'recipe': 'qwen3_8b_qlora_oracle_sft',\n",
441
+ " 'base_model': MODEL,\n",
442
+ " 'artifact_repo': ARTIFACT_REPO,\n",
443
+ " 'space_repo': SPACE_REPO,\n",
444
+ " 'sft_case_count': SFT_DATASET_SIZE,\n",
445
+ " 'sft_row_count': len(dataset),\n",
446
+ " 'sft_epochs': SFT_EPOCHS,\n",
447
+ " 'sft_max_steps': SFT_MAX_STEPS,\n",
448
+ " 'sft_learning_rate': SFT_LEARNING_RATE,\n",
449
+ " 'max_sft_length': MAX_SFT_LENGTH,\n",
450
+ " 'lora_r': LORA_R,\n",
451
+ " 'lora_alpha': LORA_ALPHA,\n",
452
+ " 'lora_dropout': LORA_DROPOUT,\n",
453
+ " 'gradient_accumulation_steps': GRADIENT_ACCUMULATION_STEPS,\n",
454
+ " 'include_rest_rows': INCLUDE_REST_ROWS,\n",
455
+ " 'metrics': getattr(train_result, 'metrics', {}) or {},\n",
456
+ "}\n",
457
+ "summary_path = OUTPUT_DIR / 'training_summary.json'\n",
458
+ "summary_path.write_text(json.dumps(summary, indent=2, sort_keys=True), encoding='utf-8')\n",
459
+ "print(json.dumps(summary, indent=2, sort_keys=True))\n",
460
+ "\n",
461
+ "token = os.getenv('HF_TOKEN')\n",
462
+ "if token:\n",
463
+ " api = HfApi(token=token)\n",
464
+ " api.create_repo(repo_id=ARTIFACT_REPO, repo_type='model', exist_ok=True)\n",
465
+ " api.upload_folder(\n",
466
+ " repo_id=ARTIFACT_REPO,\n",
467
+ " repo_type='model',\n",
468
+ " folder_path=str(OUTPUT_DIR),\n",
469
+ " commit_message=f'Upload Counsel-Env QLoRA SFT adapter ({SFT_MAX_STEPS} max steps)',\n",
470
+ " )\n",
471
+ " print(f'Uploaded QLoRA SFT adapter to https://huggingface.co/{ARTIFACT_REPO}')\n",
472
+ "else:\n",
473
+ " print('HF_TOKEN not set; skipping artifact upload.')"
474
+ ]
475
+ },
476
+ {
477
+ "cell_type": "markdown",
478
+ "metadata": {},
479
+ "source": [
480
+ "## 10. Reproducing run4b numbers\n",
481
+ "\n",
482
+ "After training finishes, evaluate against the same 150 deterministic seeds used for the official scorecard:\n",
483
+ "\n",
484
+ "```\n",
485
+ "python counsel_env/scripts/evaluate_trained_checkpoint_job.py\n",
486
+ "```\n",
487
+ "\n",
488
+ "Expected on run4b:\n",
489
+ "\n",
490
+ "| Metric | Value |\n",
491
+ "| --- | ---: |\n",
492
+ "| `train_loss` | 0.0565 |\n",
493
+ "| `train_runtime` (s) | 1287.7 |\n",
494
+ "| `avg_reward` (150 seeds) | 0.864 |\n",
495
+ "| `primary_reward` (150 seeds) | 0.943 |\n",
496
+ "| `surface_rate` | 0.943 |\n",
497
+ "| `invalid_tool_calls` | 0 |"
498
+ ]
499
+ }
500
+ ],
501
+ "metadata": {
502
+ "kernelspec": {
503
+ "display_name": "Python 3",
504
+ "language": "python",
505
+ "name": "python3"
506
+ },
507
+ "language_info": {
508
+ "name": "python",
509
+ "version": "3.11"
510
+ }
511
+ },
512
+ "nbformat": 4,
513
+ "nbformat_minor": 5
514
+ }