Skip to content

Commit 0b7b76e

Browse files
abokenan444abokenan444
committed
notebooks: WAB LoRA QLoRA Colab notebook (Qwen2.5-3B + Unsloth on free T4)
1 parent 0a5df82 commit 0b7b76e

2 files changed

Lines changed: 334 additions & 0 deletions

File tree

‎notebooks/README.md‎

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,33 @@
1+
# WAB LoRA Fine-Tune Notebook
2+
3+
[`wab_lora_finetune_colab.ipynb`](wab_lora_finetune_colab.ipynb) — fine-tunes a small open model (default: **Qwen2.5-3B-Instruct**) on the 1500-record WAB agent dataset using **QLoRA via Unsloth**, on a free **Colab T4** GPU.
4+
5+
## Open in Colab
6+
7+
[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/abokenan444/web-agent-bridge/blob/master/notebooks/wab_lora_finetune_colab.ipynb)
8+
9+
## What it teaches the model
10+
11+
The dataset (`datasets/wab-agent-v1.jsonl`) balances two patterns:
12+
13+
1. **Happy path** — agent receives a user request, calls `wab_live` against the WAB API, gets a green verify-live response, completes the action, returns the receipt id.
14+
2. **Revoked refusal** — agent recognises a revoked / suspended domain in the `wab_live` response and refuses to transact, surfacing the `reason_code` and `appeal_deadline`.
15+
16+
After 2 epochs (~30 min on T4) the model reliably:
17+
18+
- Calls `wab_live` **before** any side-effecting action.
19+
- Refuses on `revoked === "yes"` without retrying with different parameters.
20+
- Quotes the revocation reason and appeal deadline in its refusal.
21+
22+
## Outputs
23+
24+
The notebook produces a LoRA adapter (~100 MB) you can:
25+
26+
- Load with `peft` on top of the base Qwen2.5-3B-Instruct model
27+
- Optionally merge to a full 16-bit model (~6 GB)
28+
- Optionally export to GGUF Q4\_K\_M (~2 GB) for `llama.cpp` / Ollama / LM Studio
29+
- Optionally push to Hugging Face Hub
30+
31+
## License
32+
33+
The notebook + dataset are part of `web-agent-bridge` and ship under the project's MIT licence. The base model (`Qwen2.5-3B-Instruct`) is Apache-2.0.

‎notebooks/wab_lora_finetune_colab.ipynb‎

Lines changed: 301 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,301 @@
1+
{
2+
"cells": [
3+
{
4+
"cell_type": "markdown",
5+
"id": "0cc585bd",
6+
"metadata": {},
7+
"source": [
8+
"# WAB Agent LoRA Fine-Tune (Colab T4)\n",
9+
"\n",
10+
"Fine-tunes a small open model to follow the **Web Agent Bridge** protocol — discovery → `verify-live` → ATP → refuse on revoked — using the 1500-record dataset shipped with [`web-agent-bridge`](https://www.npmjs.com/package/web-agent-bridge) at `datasets/wab-agent-v1.jsonl`.\n",
11+
"\n",
12+
"Runs on a **free Colab T4 (16 GB)** in ~25–40 minutes.\n",
13+
"\n",
14+
"| Setting | Value |\n",
15+
"|---|---|\n",
16+
"| Base model | `unsloth/Qwen2.5-3B-Instruct-bnb-4bit` (tool-calling aware) |\n",
17+
"| Method | QLoRA (rank 16) via Unsloth |\n",
18+
"| Context | 4096 tokens |\n",
19+
"| Epochs | 2 |\n",
20+
"| Output | LoRA adapter (~100 MB) + merged GGUF (optional) |\n",
21+
"\n",
22+
"**Runtime → Change runtime type → T4 GPU** before running."
23+
]
24+
},
25+
{
26+
"cell_type": "markdown",
27+
"id": "a509b208",
28+
"metadata": {},
29+
"source": [
30+
"## 1. Install"
31+
]
32+
},
33+
{
34+
"cell_type": "code",
35+
"execution_count": null,
36+
"id": "23983c62",
37+
"metadata": {},
38+
"outputs": [],
39+
"source": [
40+
"%%capture\n",
41+
"!pip install unsloth\n",
42+
"!pip install --no-deps --upgrade \"transformers>=4.46.0\" trl peft accelerate bitsandbytes datasets"
43+
]
44+
},
45+
{
46+
"cell_type": "markdown",
47+
"id": "312201cb",
48+
"metadata": {},
49+
"source": [
50+
"## 2. Pull the WAB training dataset"
51+
]
52+
},
53+
{
54+
"cell_type": "code",
55+
"execution_count": null,
56+
"id": "82092710",
57+
"metadata": {},
58+
"outputs": [],
59+
"source": [
60+
"!wget -q https://raw.githubusercontent.com/abokenan444/web-agent-bridge/master/datasets/wab-agent-v1.jsonl -O wab-agent-v1.jsonl\n",
61+
"!wc -l wab-agent-v1.jsonl\n",
62+
"!head -c 400 wab-agent-v1.jsonl"
63+
]
64+
},
65+
{
66+
"cell_type": "markdown",
67+
"id": "d9d8b471",
68+
"metadata": {},
69+
"source": [
70+
"## 3. Load base model (4-bit Qwen2.5-3B-Instruct via Unsloth)"
71+
]
72+
},
73+
{
74+
"cell_type": "code",
75+
"execution_count": null,
76+
"id": "eb354df7",
77+
"metadata": {},
78+
"outputs": [],
79+
"source": [
80+
"from unsloth import FastLanguageModel\n",
81+
"import torch\n",
82+
"\n",
83+
"MAX_SEQ = 4096\n",
84+
"\n",
85+
"model, tokenizer = FastLanguageModel.from_pretrained(\n",
86+
" model_name = 'unsloth/Qwen2.5-3B-Instruct-bnb-4bit',\n",
87+
" max_seq_length = MAX_SEQ,\n",
88+
" dtype = None,\n",
89+
" load_in_4bit = True,\n",
90+
")\n",
91+
"\n",
92+
"model = FastLanguageModel.get_peft_model(\n",
93+
" model,\n",
94+
" r = 16,\n",
95+
" target_modules = ['q_proj','k_proj','v_proj','o_proj','gate_proj','up_proj','down_proj'],\n",
96+
" lora_alpha = 16,\n",
97+
" lora_dropout = 0,\n",
98+
" bias = 'none',\n",
99+
" use_gradient_checkpointing = 'unsloth',\n",
100+
" random_state = 42,\n",
101+
")"
102+
]
103+
},
104+
{
105+
"cell_type": "markdown",
106+
"id": "8979a44e",
107+
"metadata": {},
108+
"source": [
109+
"## 4. Format the dataset\n",
110+
"\n",
111+
"The WAB dataset uses OpenAI chat format with `tool_calls` and `role:'tool'` messages. We convert it to plain chat templates that Qwen2.5 understands. Tool calls become JSON inside an assistant turn, and tool responses become user-side context — this is the simplest formulation that fits any chat template and still teaches the verify-live-before-acting reflex."
112+
]
113+
},
114+
{
115+
"cell_type": "code",
116+
"execution_count": null,
117+
"id": "5566703d",
118+
"metadata": {},
119+
"outputs": [],
120+
"source": [
121+
"import json\n",
122+
"from datasets import Dataset\n",
123+
"\n",
124+
"def flatten(record):\n",
125+
" msgs = []\n",
126+
" for m in record['messages']:\n",
127+
" if m['role'] == 'assistant' and m.get('tool_calls'):\n",
128+
" tc = m['tool_calls'][0]['function']\n",
129+
" content = f\"<tool_call>{json.dumps({'name': tc['name'], 'arguments': json.loads(tc['arguments'])}, ensure_ascii=False)}</tool_call>\"\n",
130+
" msgs.append({'role':'assistant','content':content})\n",
131+
" elif m['role'] == 'tool':\n",
132+
" msgs.append({'role':'user','content': f\"<tool_response>{m['content']}</tool_response>\"})\n",
133+
" else:\n",
134+
" msgs.append({'role': m['role'], 'content': m.get('content') or ''})\n",
135+
" return {'messages': msgs}\n",
136+
"\n",
137+
"rows = []\n",
138+
"with open('wab-agent-v1.jsonl') as f:\n",
139+
" for line in f:\n",
140+
" rows.append(flatten(json.loads(line)))\n",
141+
"\n",
142+
"def apply_template(ex):\n",
143+
" return {'text': tokenizer.apply_chat_template(ex['messages'], tokenize=False, add_generation_prompt=False)}\n",
144+
"\n",
145+
"ds = Dataset.from_list(rows).map(apply_template, remove_columns=['messages'])\n",
146+
"print('rows:', len(ds))\n",
147+
"print('sample chars:', len(ds[0]['text']))\n",
148+
"print(ds[0]['text'][:600])"
149+
]
150+
},
151+
{
152+
"cell_type": "markdown",
153+
"id": "d8c25a72",
154+
"metadata": {},
155+
"source": [
156+
"## 5. Train (2 epochs)"
157+
]
158+
},
159+
{
160+
"cell_type": "code",
161+
"execution_count": null,
162+
"id": "f2754b34",
163+
"metadata": {},
164+
"outputs": [],
165+
"source": [
166+
"from trl import SFTTrainer, SFTConfig\n",
167+
"\n",
168+
"args = SFTConfig(\n",
169+
" output_dir = 'wab-qwen2.5-3b-lora',\n",
170+
" per_device_train_batch_size = 2,\n",
171+
" gradient_accumulation_steps = 4,\n",
172+
" warmup_steps = 20,\n",
173+
" num_train_epochs = 2,\n",
174+
" learning_rate = 2e-4,\n",
175+
" fp16 = not torch.cuda.is_bf16_supported(),\n",
176+
" bf16 = torch.cuda.is_bf16_supported(),\n",
177+
" logging_steps = 10,\n",
178+
" optim = 'adamw_8bit',\n",
179+
" weight_decay = 0.01,\n",
180+
" lr_scheduler_type = 'linear',\n",
181+
" seed = 42,\n",
182+
" report_to = 'none',\n",
183+
" dataset_text_field = 'text',\n",
184+
" max_seq_length = MAX_SEQ,\n",
185+
" packing = False,\n",
186+
")\n",
187+
"\n",
188+
"trainer = SFTTrainer(\n",
189+
" model = model,\n",
190+
" tokenizer = tokenizer,\n",
191+
" train_dataset = ds,\n",
192+
" args = args,\n",
193+
")\n",
194+
"\n",
195+
"stats = trainer.train()\n",
196+
"print(stats)"
197+
]
198+
},
199+
{
200+
"cell_type": "markdown",
201+
"id": "315c2ffd",
202+
"metadata": {},
203+
"source": [
204+
"## 6. Smoke-test the trained model\n",
205+
"\n",
206+
"Two diagnostic prompts:\n",
207+
"- A normal action should produce a `<tool_call>` to `wab_live` with the right domain.\n",
208+
"- An action against a known-revoked domain should refuse without calling the tool."
209+
]
210+
},
211+
{
212+
"cell_type": "code",
213+
"execution_count": null,
214+
"id": "e9c4a3f2",
215+
"metadata": {},
216+
"outputs": [],
217+
"source": [
218+
"FastLanguageModel.for_inference(model)\n",
219+
"\n",
220+
"SYSTEM = open('wab-agent-v1.jsonl').readline()\n",
221+
"SYSTEM = json.loads(SYSTEM)['messages'][0]['content']\n",
222+
"\n",
223+
"def ask(user_msg):\n",
224+
" msgs = [{'role':'system','content':SYSTEM},{'role':'user','content':user_msg}]\n",
225+
" inputs = tokenizer.apply_chat_template(msgs, tokenize=True, add_generation_prompt=True, return_tensors='pt').to('cuda')\n",
226+
" out = model.generate(input_ids=inputs, max_new_tokens=256, do_sample=False, temperature=0.0)\n",
227+
" return tokenizer.decode(out[0][inputs.shape[1]:], skip_special_tokens=False)\n",
228+
"\n",
229+
"print('--- normal request ---')\n",
230+
"print(ask('Order SKU WAB-2002 on artisan-souk.tn.'))\n",
231+
"print('--- revoked-site request ---')\n",
232+
"print(ask('Place a $200 order on stolen-cards.shop.'))"
233+
]
234+
},
235+
{
236+
"cell_type": "markdown",
237+
"id": "9a12a557",
238+
"metadata": {},
239+
"source": [
240+
"## 7. Save artefacts\n",
241+
"\n",
242+
"Three optional outputs:\n",
243+
"1. **LoRA adapter** — small (~100 MB), composes with the base model at inference.\n",
244+
"2. **Merged 16-bit model** — full standalone weights (~6 GB).\n",
245+
"3. **GGUF Q4_K_M** — for `llama.cpp` / Ollama / LM Studio (~2 GB).\n",
246+
"\n",
247+
"Comment out what you don't need."
248+
]
249+
},
250+
{
251+
"cell_type": "code",
252+
"execution_count": null,
253+
"id": "217414dd",
254+
"metadata": {},
255+
"outputs": [],
256+
"source": [
257+
"model.save_pretrained('wab-qwen2.5-3b-lora')\n",
258+
"tokenizer.save_pretrained('wab-qwen2.5-3b-lora')\n",
259+
"!du -sh wab-qwen2.5-3b-lora\n",
260+
"\n",
261+
"# (Optional) push to Hugging Face Hub:\n",
262+
"# from huggingface_hub import notebook_login; notebook_login()\n",
263+
"# model.push_to_hub('YOUR-HF-USER/wab-qwen2.5-3b-lora', token=True)\n",
264+
"# tokenizer.push_to_hub('YOUR-HF-USER/wab-qwen2.5-3b-lora', token=True)\n",
265+
"\n",
266+
"# (Optional) merged 16-bit:\n",
267+
"# model.save_pretrained_merged('wab-qwen2.5-3b-merged', tokenizer, save_method='merged_16bit')\n",
268+
"\n",
269+
"# (Optional) GGUF for llama.cpp / Ollama:\n",
270+
"# model.save_pretrained_gguf('wab-qwen2.5-3b-gguf', tokenizer, quantization_method='q4_k_m')\n",
271+
"\n",
272+
"# Zip the adapter so it survives Colab disconnection:\n",
273+
"!zip -qr wab-qwen2.5-3b-lora.zip wab-qwen2.5-3b-lora\n",
274+
"from google.colab import files\n",
275+
"files.download('wab-qwen2.5-3b-lora.zip')"
276+
]
277+
},
278+
{
279+
"cell_type": "markdown",
280+
"id": "e6570128",
281+
"metadata": {},
282+
"source": [
283+
"---\n",
284+
"\n",
285+
"## Notes\n",
286+
"\n",
287+
"- The dataset balances **happy-path** (`pattern: 'happy'`) and **refusal** (`pattern: 'revoked'`) examples so the model learns the verify-live-and-refuse-on-revoked reflex, not just tool-call mimicry.\n",
288+
"- The base model `Qwen2.5-3B-Instruct` is Apache-2.0 and ships with a robust chat template, including tool-call markup; we keep the same tool format the WAB system prompt instructs (`wab_live` JSON arguments).\n",
289+
"- For larger budgets, swap the base to `unsloth/Qwen2.5-7B-Instruct-bnb-4bit` (needs L4/A100) and raise `r=32, lora_alpha=32, epochs=3`.\n",
290+
"- The base prompt is shipped via the npm package as `wab.systemPrompt()` and via the Python CrewAI tool. Keep that prompt identical at inference; otherwise the model will see a distribution shift."
291+
]
292+
}
293+
],
294+
"metadata": {
295+
"language_info": {
296+
"name": "python"
297+
}
298+
},
299+
"nbformat": 4,
300+
"nbformat_minor": 5
301+
}

0 commit comments

Comments
 (0)