On this tutorial, we implement an end-to-end supervised fine-tuning pipeline for the XYZ-Aquila-SFT dataset, Hugging Face Transformers, PyTorch, and PEFT. We stream and examine the dataset, parse multi-turn tool-use trajectories, extract structured instrument calls, analyze corpus traits, and protect embedded reasoning and remark patterns. We then convert instrument schemas between message-embedded and structured codecs, render Qwen-compatible ChatML with assistant-only loss masking, put together a customized PyTorch dataset and collator, and fine-tune Qwen3-0.6B with LoRA. Lastly, we consider tool-call prediction earlier than and after coaching and export each the remodeled dataset and corpus statistics for additional experimentation.
import os, sys, subprocess
CFG = dict(
REPO = "XYZAILab/XYZ-Aquila-SFT",
LANG = "en",
N_STREAM = 400,
N_EVAL = 40,
MODEL_ID = "Qwen/Qwen3-0.6B",
MAX_SEQ_LEN = 2048,
LENGTH_POLICY = "truncate",
RUN_TRAINING = True,
MAX_STEPS = 30,
GRAD_ACCUM = 8,
LR = 1e-4,
LORA_R = 16,
RUN_EVAL = True,
N_EVAL_PROBES = 24,
OUT_DIR = "/content material/aquila_out",
SEED = 0,
)
os.makedirs(CFG["OUT_DIR"], exist_ok=True)
def pip(*pkgs):
subprocess.run([sys.executable, "-m", "pip", "install", "-q", "-U", *pkgs], examine=False)
pip("datasets>=3.0.0", "transformers>=4.51.0", "peft>=0.13.0", "speed up>=1.0.0")
import json, re, math, random, statistics as stats
from collections import Counter, defaultdict
from dataclasses import dataclass, discipline
from typing import Any, Dict, Record, Non-compulsory
import torch
import matplotlib.pyplot as plt
from datasets import load_dataset
from transformers import AutoTokenizer, AutoModelForCausalLM, get_cosine_schedule_with_warmup
random.seed(CFG["SEED"]); torch.manual_seed(CFG["SEED"])
DEV = "cuda" if torch.cuda.is_available() else "cpu"
BF16 = DEV == "cuda" and torch.cuda.is_bf16_supported()
print(f"system={DEV} bf16={BF16} torch={torch.__version__}")
print(f"n[1] streaming {CFG['REPO']}:{CFG['LANG']} ...")
stream = load_dataset(CFG["REPO"], CFG["LANG"], break up="practice", streaming=True)
RAW: Record[Dict[str, Any]] = record(stream.take(CFG["N_STREAM"]))
print(f" pulled {len(RAW)} rows; keys = {record(RAW[0].keys())}")
_r = RAW[0]
print(f" query[:110] : {_r['question'][:110]}...")
print(f" reply : {_r['answer'][:80]}")
print(f" variety of instrument calls : {_r['number of tool calls']}")
print(f" trajectory len : {len(_r['trajectory'])} msgs")
print(f" function sequence (first8): {[m['role'] for m in _r['trajectory'][:8]]}")
We configure the dataset, mannequin, coaching parameters, output listing, and reproducibility settings for the entire workflow. We set up the required Hugging Face, PEFT, Speed up, and PyTorch-related dependencies and detect whether or not a CUDA GPU and BF16 help can be found. We then stream a restricted variety of XYZ-Aquila-SFT examples, examine the dataset schema, and look at the construction of the primary tool-use trajectory.
TOOLS_BLOCK_RE = re.compile(r"s*(.*?)s* ", re.S)
THINK_RE = re.compile(r"(.*?) ", re.S)
TOOL_RESP_RE = re.compile(r"s*(.*?)s* ", re.S)
TOOLS_HDR_RE = re.compile(r"nn# Toolsnn")
def iter_json_objects(textual content: str, restrict: int = 1):
"""Nesting-safe JSON scanner. Regex like r'{.*?}' breaks on nested
`arguments` objects, which each and every actual instrument name has."""
dec, i, n, out = json.JSONDecoder(), 0, len(textual content), []
whereas i < n and len(out) < restrict:
whereas i < n and textual content[i] not in "{[":
i += 1
if i >= n:
break
try:
obj, end = dec.raw_decode(text, i)
except json.JSONDecodeError:
i += 1
continue
out.append(obj); i = end
return out
def parse_tool_calls(content: str) -> List[Dict[str, Any]]:
calls = []
for m in re.finditer(r"", content material):
acquired = iter_json_objects(content material[m.end():], restrict=1)
if acquired:
calls.append(acquired[0])
return calls
@dataclass
class Trajectory:
query: str
reply: str
declared_calls: int
messages: Record[Dict[str, str]]
system_core: str = ""
instruments: Record[Dict[str, Any]] = discipline(default_factory=record)
tools_suffix: str = ""
calls: Record[Dict[str, Any]] = discipline(default_factory=record)
n_observations: int = 0
n_think: int = 0
@property
def tool_names(self): return [c.get("name", "?") for c in self.calls]
@property
def depth(self): return len(self.messages)
def parse_row(row: Dict[str, Any]) -> Trajectory:
msgs = [{"role": m["role"], "content material": m["content"]} for m in row["trajectory"]]
t = Trajectory(row["question"], row["answer"], row["number of tool calls"], msgs)
if msgs and msgs[0]["role"] == "system":
sysmsg = msgs[0]["content"]
break up = TOOLS_HDR_RE.search(sysmsg)
if break up:
t.system_core = sysmsg[:split.start()]
t.tools_suffix = sysmsg[split.start():]
else:
t.system_core = sysmsg
blk = TOOLS_BLOCK_RE.search(sysmsg)
if blk:
t.instruments = iter_json_objects(blk.group(1), restrict=64)
for m in msgs:
if m["role"] == "assistant":
t.calls += parse_tool_calls(m["content"])
t.n_think += len(THINK_RE.findall(m["content"]))
else:
t.n_observations += len(TOOL_RESP_RE.findall(m["content"]))
return t
TRAJ = [parse_row(r) for r in RAW]
t0 = TRAJ[0]
print(f"n[2] parsed {len(TRAJ)} trajectories")
print(f" instrument schemas discovered : {[fn.get('function', fn).get('name') for fn in t0.tools]}")
print(f" parsed calls : {len(t0.calls)} (declared {t0.declared_calls})")
print(f" observations : {t0.n_observations} suppose blocks: {t0.n_think}")
if t0.calls:
print(f" pattern name : {json.dumps(t0.calls[0], ensure_ascii=False)[:200]}")
agree = sum(len(t.calls) == t.declared_calls for t in TRAJ)
print(f" parser vs 'variety of instrument calls': {agree}/{len(TRAJ)} precise match")
calls_per = [len(t.calls) for t in TRAJ]
depth_per = [t.depth for t in TRAJ]
chars_per = [sum(len(m["content"]) for m in t.messages) for t in TRAJ]
name_freq = Counter(n for t in TRAJ for n in t.tool_names)
argkey_freq = defaultdict(Counter)
for t in TRAJ:
for c in t.calls:
args = c.get("arguments", {})
if isinstance(args, dict):
for okay in args: argkey_freq[c.get("name", "?")][k] += 1
def q(xs, p):
xs = sorted(xs); return xs[min(len(xs) - 1, int(p * len(xs)))]
print("n[3] corpus statistics")
print(f" instrument calls / traj : imply {stats.imply(calls_per):.1f} p50 {q(calls_per,.5)} "
f"p90 {q(calls_per,.9)} max {max(calls_per)}")
print(f" messages / traj : imply {stats.imply(depth_per):.1f} p90 {q(depth_per,.9)} max {max(depth_per)}")
print(f" chars / traj : imply {stats.imply(chars_per):,.0f} p90 {q(chars_per,.9):,}")
print(f" instrument distribution : {dict(name_freq)}")
for okay, v in argkey_freq.gadgets():
print(f" {okay:<24} arg keys -> {dict(v.most_common(6))}")
tot = sum(chars_per); high = sum(sorted(chars_per)[-max(1, len(chars_per)//10):])
print(f" top-10% longest trajectories maintain {100*high/tot:.1f}% of all characters")
fig, ax = plt.subplots(1, 3, figsize=(15, 3.6))
ax[0].hist(calls_per, bins=40); ax[0].set_yscale("log"); ax[0].set_title("instrument calls / trajectory")
ax[1].hist(depth_per, bins=40); ax[1].set_yscale("log"); ax[1].set_title("messages / trajectory")
ax[2].bar(record(name_freq), record(name_freq.values())); ax[2].set_title("instrument utilization"); ax[2].tick_params(axis="x", rotation=20)
plt.tight_layout(); plt.present()
We outline nesting-safe utilities for extracting JSON instrument calls, reasoning blocks, observations, and embedded instrument schemas from every dialog. We convert each uncooked dataset row right into a structured trajectory object and confirm that the parsed tool-call counts match the values declared by the dataset. We then calculate corpus-level statistics and visualize the distributions of instrument calls, message depth, trajectory measurement, and power utilization frequency.
QWEN3_TOOLS_TMPL = (
"You might be supplied with perform signatures inside n"
"{strains}n nnFor every perform name, return a json object with perform title "
"and arguments inside n"
'{{"title": , "arguments": }}n '
)
def extract_tools(t: Trajectory) -> Dict[str, Any]:
"""message-embedded schemas -> {'messages': [...], 'instruments': [...]}"""
msgs = [dict(m) for m in t.messages]
if msgs and msgs[0]["role"] == "system":
msgs[0]["content"] = t.system_core
return {"messages": msgs, "instruments": t.instruments,
"query": t.query, "reply": t.reply}
def render_tools(rec: Dict[str, Any]) -> Record[Dict[str, str]]:
"""inverse: structured instruments -> schemas re-embedded within the system message"""
msgs = [dict(m) for m in rec["messages"]]
if rec["tools"] and msgs and msgs[0]["role"] == "system":
strains = "n".be a part of(json.dumps(x, ensure_ascii=False) for x in rec["tools"])
msgs[0]["content"] = msgs[0]["content"] + QWEN3_TOOLS_TMPL.format(strains=strains)
return msgs
_rt = render_tools(extract_tools(t0))
precise = _rt[0]["content"] == t0.messages[0]["content"]
print(f"n[4] extract->render byte-exact: {precise}")
if not precise:
print(" template drift detected -> utilizing verbatim tools_suffix for render()")
a, b = t0.messages[0]["content"], _rt[0]["content"]
i = subsequent((i for i in vary(min(len(a), len(b))) if a[i] != b[i]), min(len(a), len(b)))
print(f" first divergence @{i}: {a[i:i+70]!r} vs {b[i:i+70]!r}")
tok = AutoTokenizer.from_pretrained(CFG["MODEL_ID"])
if tok.pad_token is None:
tok.pad_token = tok.eos_token
IM_START, IM_END, NL = "<|im_start|>", "<|im_end|>", "n"
def render_and_mask(t: Trajectory, max_len: int, coverage: str):
"""Handbook ChatML so we management masking token-exactly.
WHY NOT apply_chat_template(): Qwen3's template deletes ...
from each assistant flip besides the final. On this dataset that silently
destroys a lot of the reasoning supervision you're paying to coach on.
"""
ids, labels = [], []
for m in t.messages:
head = tok(f"{IM_START}{m['role']}{NL}", add_special_tokens=False).input_ids
physique = tok(m["content"], add_special_tokens=False).input_ids
tail = tok(f"{IM_END}{NL}", add_special_tokens=False).input_ids
seg = head + physique + tail
if m["role"] == "assistant":
lab = [-100] * len(head) + physique + tail
else:
lab = [-100] * len(seg)
ids += seg; labels += lab
if len(ids) > max_len:
if coverage == "drop":
return None
ids, labels = ids[:max_len], labels[:max_len]
if all(l == -100 for l in labels):
return None
return {"input_ids": ids, "labels": labels}
_probe = [{"role": "system", "content": "S"}, {"role": "user", "content": "U"},
{"role": "assistant", "content": "A"}]
_mine = "".be a part of(f"{IM_START}{m['role']}{NL}{m['content']}{IM_END}{NL}" for m in _probe)
_theirs = tok.apply_chat_template(_probe, tokenize=False, add_generation_prompt=False)
print(f"n[5] handbook ChatML == chat_template on tool-free probe: {_mine == _theirs}")
if _mine != _theirs:
print(f" mine : {_mine!r}n theirs: {_theirs!r} (informational solely)")
ENC = [e for e in (render_and_mask(t, CFG["MAX_SEQ_LEN"], CFG["LENGTH_POLICY"]) for t in TRAJ) if e]
sup = [sum(1 for x in e["labels"] if x != -100) / len(e["labels"]) for e in ENC]
print(f" encoded {len(ENC)}/{len(TRAJ)} examples")
print(f" supervised-token ratio: imply {stats.imply(sup):.3f} p10 {q(sup,.1):.3f} p90 {q(sup,.9):.3f}")
over = sum(1 for t in TRAJ if sum(len(tok(m['content'], add_special_tokens=False).input_ids)
for m in t.messages[:3]) > CFG["MAX_SEQ_LEN"])
print(f" trajectories whose first 3 msgs alone exceed MAX_SEQ_LEN: {over}")
SPLIT = len(ENC) - min(CFG["N_EVAL"], len(ENC)//5)
TRAIN_ENC, EVAL_TRAJ = ENC[:SPLIT], TRAJ[SPLIT:]
class SFTSet(torch.utils.information.Dataset):
def __init__(self, rows): self.rows = rows
def __len__(self): return len(self.rows)
def __getitem__(self, i): return self.rows[i]
def collate(batch):
L = max(len(b["input_ids"]) for b in batch)
pad = tok.pad_token_id
return {
"input_ids": torch.tensor([b["input_ids"] + [pad]*(L-len(b["input_ids"])) for b in batch]),
"labels": torch.tensor([b["labels"] + [-100]*(L-len(b["labels"])) for b in batch]),
"attention_mask": torch.tensor([[1]*len(b["input_ids"]) + [0]*(L-len(b["input_ids"])) for b in batch]),
}
loader = torch.utils.information.DataLoader(SFTSet(TRAIN_ENC), batch_size=1, shuffle=True, collate_fn=collate)
print(f"n[6] practice={len(TRAIN_ENC)} eval_trajectories={len(EVAL_TRAJ)}")
We extract embedded instrument definitions right into a structured format and reconstruct them to check whether or not the conversion preserves the unique system message. We manually render every trajectory in ChatML format to retain all reasoning content material and apply loss solely to assistant-generated tokens. We additionally tokenize the examples, implement the chosen sequence-length coverage, create the coaching and analysis break up, and put together a padded PyTorch DataLoader.
def build_probes(trajs, n):
"""Instructor-forced probes: reduce the trajectory proper earlier than an assistant flip
that points a instrument name; the gold label is that decision."""
probes = []
for t in trajs:
for i, m in enumerate(t.messages):
if m["role"] != "assistant":
proceed
gold = parse_tool_calls(m["content"])
if not gold:
proceed
prefix = "".be a part of(f"{IM_START}x['role']{NL}" for x in [])
prefix = "".be a part of(f"{IM_START}{p['role']}{NL}{p['content']}{IM_END}{NL}"
for p in t.messages[:i]) + f"{IM_START}assistant{NL}"
if len(tok(prefix, add_special_tokens=False).input_ids) > CFG["MAX_SEQ_LEN"] - 160:
proceed
probes.append({"prefix": prefix, "gold": gold[0]})
break
if len(probes) >= n:
break
return probes
@torch.no_grad()
def eval_tool_calls(mannequin, probes, tag):
mannequin.eval()
name_hit = arg_f1 = parsed = 0
for p in probes:
enc = tok(p["prefix"], return_tensors="pt", add_special_tokens=False).to(mannequin.system)
out = mannequin.generate(**enc, max_new_tokens=160, do_sample=False,
pad_token_id=tok.pad_token_id)
gen = tok.decode(out[0][enc.input_ids.shape[1]:], skip_special_tokens=True)
pred = (parse_tool_calls(gen) or iter_json_objects(gen, restrict=1) or [None])[0]
if not isinstance(pred, dict):
proceed
parsed += 1
g = p["gold"]
name_hit += int(pred.get("title") == g.get("title"))
pk = set((pred.get("arguments") or {}).keys()) if isinstance(pred.get("arguments"), dict) else set()
gk = set((g.get("arguments") or {}).keys()) if isinstance(g.get("arguments"), dict) else set()
if pk or gk:
inter = len(pk & gk)
arg_f1 += 0.0 if inter == 0 else 2*inter/(len(pk)+len(gk))
n = max(1, len(probes))
print(f" [{tag}] parseable {parsed}/{n} | tool-name acc {name_hit/n:.3f} | arg-key F1 {arg_f1/n:.3f}")
return dict(parsed=parsed/n, name_acc=name_hit/n, arg_f1=arg_f1/n)
PROBES = build_probes(EVAL_TRAJ, CFG["N_EVAL_PROBES"])
print(f" constructed {len(PROBES)} teacher-forced probes")
outcomes = {}
if CFG["RUN_TRAINING"]:
from peft import LoraConfig, get_peft_model
dtype = torch.bfloat16 if BF16 else torch.float32
mannequin = AutoModelForCausalLM.from_pretrained(
CFG["MODEL_ID"], torch_dtype=dtype, attn_implementation="sdpa").to(DEV)
mannequin.config.use_cache = False
mannequin.gradient_checkpointing_enable()
mannequin.enable_input_require_grads()
if CFG["RUN_EVAL"] and PROBES and DEV == "cuda":
print("n[8] baseline eval")
outcomes["before"] = eval_tool_calls(mannequin, PROBES, "base")
mannequin = get_peft_model(mannequin, LoraConfig(
r=CFG["LORA_R"], lora_alpha=2*CFG["LORA_R"], lora_dropout=0.05,
bias="none", task_type="CAUSAL_LM",
mannequin.print_trainable_parameters()
decide = torch.optim.AdamW([p for p in model.parameters() if p.requires_grad],
lr=CFG["LR"], weight_decay=0.0, betas=(0.9, 0.95))
sched = get_cosine_schedule_with_warmup(decide, 5, CFG["MAX_STEPS"])
scaler = torch.amp.GradScaler("cuda", enabled=(DEV == "cuda" and never BF16))
amp_dt = torch.bfloat16 if BF16 else torch.float16
print(f"n[7] coaching {CFG['MAX_STEPS']} steps "
f"(bs1 x accum{CFG['GRAD_ACCUM']} = {CFG['GRAD_ACCUM']} traj/step)")
mannequin.practice(); step = 0; run = None; it = iter(loader)
whereas step < CFG["MAX_STEPS"]:
decide.zero_grad(set_to_none=True); acc = 0.0
for _ in vary(CFG["GRAD_ACCUM"]):
strive: batch = subsequent(it)
besides StopIteration:
it = iter(loader); batch = subsequent(it)
batch = {okay: v.to(DEV) for okay, v in batch.gadgets()}
with torch.autocast(DEV, dtype=amp_dt, enabled=(DEV == "cuda")):
loss = mannequin(**batch).loss / CFG["GRAD_ACCUM"]
scaler.scale(loss).backward() if scaler.is_enabled() else loss.backward()
acc += loss.merchandise()
if scaler.is_enabled():
scaler.unscale_(decide)
(scaler.step(decide), scaler.replace()) if scaler.is_enabled() else decide.step()
sched.step(); step += 1
run = acc if run is None else 0.9*run + 0.1*acc
if step % 5 == 0 or step == 1:
print(f" step {step:>3}/{CFG['MAX_STEPS']} loss {acc:.4f} ema {run:.4f} "
f"lr {sched.get_last_lr()[0]:.2e} ppl {math.exp(min(20, acc)):.1f}")
mannequin.save_pretrained(f"{CFG['OUT_DIR']}/lora_adapter"); tok.save_pretrained(f"{CFG['OUT_DIR']}/lora_adapter")
print(f" adapter -> {CFG['OUT_DIR']}/lora_adapter")
if CFG["RUN_EVAL"] and PROBES and DEV == "cuda":
print("n[8] post-training eval")
mannequin.config.use_cache = True
outcomes["after"] = eval_tool_calls(mannequin, PROBES, "lora")
mannequin.config.use_cache = False
if "earlier than" in outcomes and "after" in outcomes:
print("n delta:", {okay: spherical(outcomes['after'][k] - outcomes['before'][k], 3)
for okay in outcomes['after']})
print(" (30 steps on ~350 trajectories is a smoke check, not a end result — "
"count on noise, and scale N_STREAM/MAX_STEPS for something actual.)")
We construct teacher-forced analysis probes by chopping trajectories instantly earlier than assistant turns that include instrument calls. We load Qwen3-0.6B, measure its baseline tool-call efficiency, connect LoRA adapters, and fine-tune the mannequin utilizing gradient accumulation, combined precision, checkpointing, clipping, and cosine learning-rate scheduling. We then consider the tailored mannequin, evaluate its metrics with the baseline, and save the skilled LoRA adapter and tokenizer.
struct_path = f"{CFG['OUT_DIR']}/aquila_{CFG['LANG']}_structured_tools.jsonl"
with open(struct_path, "w", encoding="utf-8") as f:
for t in TRAJ:
f.write(json.dumps(extract_tools(t), ensure_ascii=False) + "n")
stats_path = f"{CFG['OUT_DIR']}/corpus_stats.json"
with open(stats_path, "w") as f:
json.dump({"n": len(TRAJ), "tool_freq": dict(name_freq),
"calls_mean": stats.imply(calls_per), "calls_max": max(calls_per),
"depth_p90": q(depth_per, .9), "encoded": len(ENC),
"supervised_ratio_mean": stats.imply(sup), "eval": outcomes}, f, indent=2)
print(f"n[9] wrote:n {struct_path}n {stats_path}")
print("performed.")
We export each parsed trajectory as a structured JSONL document containing messages, instrument schemas, questions, and solutions. We additionally save a JSON report containing corpus measurement, instrument frequencies, trajectory statistics, supervised-token ratios, and obtainable analysis outcomes. We end the workflow with reusable dataset artifacts, analytical outputs, and mannequin information saved within the configured output listing.
In conclusion, we accomplished a sensible pipeline for analyzing, reworking, fine-tuning, and evaluating complicated tool-use trajectories from the XYZ-Aquila-SFT dataset. We preserved the unique conversational construction, utilized token-level supervision solely to assistant responses, and used LoRA to adapt Qwen3-0.6B effectively on a Colab-compatible GPU. We additionally in contrast baseline and post-training tool-call efficiency via teacher-forced analysis and exported reusable structured information, mannequin adapters, and analytical statistics. This workflow offers us a robust basis for scaling tool-aware supervised fine-tuning, testing different sequence-length insurance policies, and coaching extra succesful agentic language fashions.
Try the FULL CODES right here. Additionally, be at liberty to comply with us on Twitter and don’t overlook to hitch our 150k+ML SubReddit and Subscribe to our E-newsletter. Wait! are you on telegram? now you possibly can be a part of us on telegram as nicely.
Have to accomplice with us for selling your GitHub Repo OR Hugging Face Web page OR Product Launch OR Webinar and so on.? Join with us
Sana Hassan, a consulting intern at Marktechpost and dual-degree scholar at IIT Madras, is obsessed with making use of expertise and AI to handle real-world challenges. With a eager curiosity in fixing sensible issues, he brings a recent perspective to the intersection of AI and real-life options.
