Дообучение LLM с вызовом инструментов: полное руководство с XYZ-Aquila-SFT и Qwen3
В этом руководстве мы реализуем сквозной конвейер контролируемой тонкой настройки для датасета XYZ-Aquila-SFT, Hugging Face Transformers, PyTorch и PEFT. Мы выполняем потоковую загрузку и исследование датасета, разбираем многошаговые траектории использования инструментов, извлекаем структурированные вызовы инструментов, анализируем характеристики корпуса и сохраняем встроенные шаблоны рассуждений и наблюдений. Затем мы преобразуем схемы инструментов между форматами, встроенными в сообщения, и структурированными форматами, формируем совместимый с Qwen ChatML с маскированием функции потерь только для ответов ассистента, подготавливаем собственные датасет PyTorch и collator, а также выполняем тонкую настройку Qwen3-0.6B с помощью LoRA. Наконец, мы оцениваем предсказание вызовов инструментов до и после обучения и экспортируем как преобразованный датасет, так и статистику корпуса для дальнейших экспериментов.
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/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], check=False)
pip("datasets>=3.0.0", "transformers>=4.51.0", "peft>=0.13.0", "accelerate>=1.0.0")
import json, re, math, random, statistics as stats
from collections import Counter, defaultdict
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional
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"device={DEV} bf16={BF16} torch={torch.__version__}")
print(f"\n[1] streaming {CFG['REPO']}:{CFG['LANG']} ...")
stream = load_dataset(CFG["REPO"], CFG["LANG"], split="train", streaming=True)
RAW: List[Dict[str, Any]] = list(stream.take(CFG["N_STREAM"]))
print(f" pulled {len(RAW)} rows; keys = {list(RAW[0].keys())}")
_r = RAW[0]
print(f" question[:110] : {_r['question'][:110]}...")
print(f" answer : {_r['answer'][:80]}")
print(f" number of tool calls : {_r['number of tool calls']}")
print(f" trajectory len : {len(_r['trajectory'])} msgs")
print(f" role sequence (first8): {[m['role'] for m in _r['trajectory'][:8]]}")
Мы настраиваем датасет, модель, параметры обучения, выходной каталог и параметры воспроизводимости для полного рабочего процесса. Мы устанавливаем необходимые зависимости Hugging Face, PEFT, Accelerate и связанные с PyTorch компоненты, а также определяем, доступны ли графический процессор CUDA и поддержка BF16. Затем мы выполняем потоковую загрузку ограниченного числа примеров XYZ-Aquila-SFT, исследуем схему датасета и изучаем структуру первой траектории использования инструментов.
TOOLS_BLOCK_RE = re.compile(r"<tools>\s*(.*?)\s*</tools>", re.S)
THINK_RE = re.compile(r"<think>(.*?)</think>", re.S)
TOOL_RESP_RE = re.compile(r"<tool_response>\s*(.*?)\s*</tool_response>", re.S)
TOOLS_HDR_RE = re.compile(r"\n\n# Tools\n\n")
def iter_json_objects(text: str, limit: int = 1):
"""Nesting-safe JSON scanner. Regex like r'\\{.*?\\}' breaks on nested
`arguments` objects, which every real tool call has."""
dec, i, n, out = json.JSONDecoder(), 0, len(text), []
while i < n and len(out) < limit:
while i < n and text[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"<tool_call>", content):
got = iter_json_objects(content[m.end():], limit=1)
if got:
calls.append(got[0])
return calls
@dataclass
class Trajectory:
question: str
answer: str
declared_calls: int
messages: List[Dict[str, str]]
system_core: str = ""
tools: List[Dict[str, Any]] = field(default_factory=list)
tools_suffix: str = ""
calls: List[Dict[str, Any]] = field(default_factory=list)
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": 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"]
split = TOOLS_HDR_RE.search(sysmsg)
if split:
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.tools = iter_json_objects(blk.group(1), limit=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" tool schemas found : {[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} think blocks: {t0.n_think}")
if t0.calls:
print(f" sample call : {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 'number of tool calls': {agree}/{len(TRAJ)} exact 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 k 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" tool calls / traj : mean {stats.mean(calls_per):.1f} p50 {q(calls_per,.5)} "
f"p90 {q(calls_per,.9)} max {max(calls_per)}")
print(f" messages / traj : mean {stats.mean(depth_per):.1f} p90 {q(depth_per,.9)} max {max(depth_per)}")
print(f" chars / traj : mean {stats.mean(chars_per):,.0f} p90 {q(chars_per,.9):,}")
print(f" tool distribution : {dict(name_freq)}")
for k, v in argkey_freq.items():
print(f" {k:<24} arg keys -> {dict(v.most_common(6))}")
tot = sum(chars_per); top = sum(sorted(chars_per)[-max(1, len(chars_per)//10):])
print(f" top-10% longest trajectories hold {100*top/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("tool calls / trajectory")
ax[1].hist(depth_per, bins=40); ax[1].set_yscale("log"); ax[1].set_title("messages / trajectory")
ax[2].bar(list(name_freq), list(name_freq.values())); ax[2].set_title("tool usage"); ax[2].tick_params(axis="x", rotation=20)
plt.tight_layout(); plt.show()
Мы определяем безопасные для вложенности утилиты для извлечения JSON-вызовов инструментов, блоков рассуждений, наблюдений и встроенных схем инструментов из каждого диалога. Мы преобразуем каждую необработанную строку датасета в структурированный объект траектории и проверяем, совпадает ли количество разобранных вызовов инструментов со значениями, указанными в датасете. Затем мы рассчитываем статистику на уровне корпуса и визуализируем распределения вызовов инструментов, глубины сообщений, размера траекторий и частоты использования инструментов.
QWEN3_TOOLS_TMPL = (
"You are provided with function signatures within <tools></tools> XML tags:\n<tools>\n"
"{lines}\n</tools>\n\nFor each function call, return a json object with function name "
"and arguments within <tool_call></tool_call> XML tags:\n<tool_call>\n"
'{{"name": <function-name>, "arguments": <args-json-object>}}\n</tool_call>'
)
def extract_tools(t: Trajectory) -> Dict[str, Any]:
"""message-embedded schemas -> {'messages': [...], 'tools': [...]}"""
msgs = [dict(m) for m in t.messages]
if msgs and msgs[0]["role"] == "system":
msgs[0]["content"] = t.system_core
return {"messages": msgs, "tools": t.tools,
"question": t.question, "answer": t.answer}
def render_tools(rec: Dict[str, Any]) -> List[Dict[str, str]]:
"""inverse: structured tools -> schemas re-embedded in the system message"""
msgs = [dict(m) for m in rec["messages"]]
if rec["tools"] and msgs and msgs[0]["role"] == "system":
lines = "\n".join(json.dumps(x, ensure_ascii=False) for x in rec["tools"])
msgs[0]["content"] = msgs[0]["content"] + QWEN3_TOOLS_TMPL.format(lines=lines)
return msgs
_rt = render_tools(extract_tools(t0))
exact = _rt[0]["content"] == t0.messages[0]["content"]
print(f"\n[4] extract->render byte-exact: {exact}")
if not exact:
print(" template drift detected -> using verbatim tools_suffix for render()")
a, b = t0.messages[0]["content"], _rt[0]["content"]
i = next((i for i in range(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, policy: str):
"""Manual ChatML so we control masking token-exactly.
WHY NOT apply_chat_template(): Qwen3's template deletes <think>...</think>
from every assistant turn except the last. On this dataset that silently
destroys most of the reasoning supervision you are paying to train on.
"""
ids, labels = [], []
for m in t.messages:
head = tok(f"{IM_START}{m['role']}{NL}", add_special_tokens=False).input_ids
body = tok(m["content"], add_special_tokens=False).input_ids
tail = tok(f"{IM_END}{NL}", add_special_tokens=False).input_ids
seg = head + body + tail
if m["role"] == "assistant":
lab = [-100] * len(head) + body + tail
else:
lab = [-100] * len(seg)
ids += seg; labels += lab
if len(ids) > max_len:
if policy == "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 = "".join(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] manual ChatML == chat_template on tool-free probe: {_mine == _theirs}")
if _mine != _theirs:
print(f" mine : {_mine!r}\n theirs: {_theirs!r} (informational only)")
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: mean {stats.mean(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.data.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.data.DataLoader(SFTSet(TRAIN_ENC), batch_size=1, shuffle=True, collate_fn=collate)
print(f"\n[6] train={len(TRAIN_ENC)} eval_trajectories={len(EVAL_TRAJ)}")
Мы извлекаем встроенные определения инструментов в структурированный формат и восстанавливаем их, чтобы проверить, сохраняет ли преобразование исходное системное сообщение. Мы вручную формируем каждую траекторию в формате ChatML, сохраняя всё содержимое рассуждений, и применяем функцию потерь только к токенам, сгенерированным ассистентом. Мы также токенизируем примеры, применяем выбранную политику ограничения длины последовательности, создаём разделение на обучающую и оценочную выборки и подготавливаем дополненный нулями DataLoader PyTorch.
def build_probes(trajs, n):
"""Teacher-forced probes: cut the trajectory right before an assistant turn
that issues a tool call; the gold label is that call."""
probes = []
for t in trajs:
for i, m in enumerate(t.messages):
if m["role"] != "assistant":
continue
gold = parse_tool_calls(m["content"])
if not gold:
continue
prefix = "".join(f"{IM_START}x['role']{NL}" for x in [])
prefix = "".join(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:
continue
probes.append({"prefix": prefix, "gold": gold[0]})
break
if len(probes) >= n:
break
return probes
@torch.no_grad()
def eval_tool_calls(model, probes, tag):
model.eval()
name_hit = arg_f1 = parsed = 0
for p in probes:
enc = tok(p["prefix"], return_tensors="pt", add_special_tokens=False).to(model.device)
out = model.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, limit=1) or [None])[0]
if not isinstance(pred, dict):
continue
parsed += 1
g = p["gold"]
name_hit += int(pred.get("name") == g.get("name"))
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" built {len(PROBES)} teacher-forced probes")
results = {}
if CFG["RUN_TRAINING"]:
from peft import LoraConfig, get_peft_model
dtype = torch.bfloat16 if BF16 else torch.float32
model = AutoModelForCausalLM.from_pretrained(
CFG["MODEL_ID"], torch_dtype=dtype, attn_implementation="sdpa").to(DEV)
model.config.use_cache = False
model.gradient_checkpointing_enable()
model.enable_input_require_grads()
if CFG["RUN_EVAL"] and PROBES and DEV == "cuda":
print("\n[8] baseline eval")
results["before"] = eval_tool_calls(model, PROBES, "base")
model = get_peft_model(model, LoraConfig(
r=CFG["LORA_R"], lora_alpha=2*CFG["LORA_R"], lora_dropout=0.05,
bias="none", task_type="CAUSAL_LM",
model.print_trainable_parameters()
opt = 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(opt, 5, CFG["MAX_STEPS"])
scaler = torch.amp.GradScaler("cuda", enabled=(DEV == "cuda" and not BF16))
amp_dt = torch.bfloat16 if BF16 else torch.float16
print(f"\n[7] training {CFG['MAX_STEPS']} steps "
f"(bs1 x accum{CFG['GRAD_ACCUM']} = {CFG['GRAD_ACCUM']} traj/step)")
model.train(); step = 0; run = None; it = iter(loader)
while step < CFG["MAX_STEPS"]:
opt.zero_grad(set_to_none=True); acc = 0.0
for _ in range(CFG["GRAD_ACCUM"]):
try: batch = next(it)
except StopIteration:
it = iter(loader); batch = next(it)
batch = {k: v.to(DEV) for k, v in batch.items()}
with torch.autocast(DEV, dtype=amp_dt, enabled=(DEV == "cuda")):
loss = model(**batch).loss / CFG["GRAD_ACCUM"]
scaler.scale(loss).backward() if scaler.is_enabled() else loss.backward()
acc += loss.item()
if scaler.is_enabled():
scaler.unscale_(opt)
(scaler.step(opt), scaler.update()) if scaler.is_enabled() else opt.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}")
model.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")
model.config.use_cache = True
results["after"] = eval_tool_calls(model, PROBES, "lora")
model.config.use_cache = False
if "before" in results and "after" in results:
print("\n delta:", {k: round(results['after'][k] - results['before'][k], 3)
for k in results['after']})
print(" (30 steps on ~350 trajectories is a smoke test, not a result — "
"expect noise, and scale N_STREAM/MAX_STEPS for anything real.)")
Мы создаём оценочные примеры с принудительной подачей правильных ответов, обрезая траектории непосредственно перед ходами ассистента, содержащими вызовы инструментов. Мы загружаем Qwen3-0.6B, измеряем исходные показатели вызова инструментов, подключаем адаптеры LoRA и выполняем тонкую настройку модели с использованием накопления градиентов, смешанной точности, контрольных точек, отсечения градиентов и косинусного планировщика скорости обучения. Затем мы оцениваем адаптированную модель, сравниваем её метрики с базовыми и сохраняем обученный адаптер LoRA и токенизатор.
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.mean(calls_per), "calls_max": max(calls_per),
"depth_p90": q(depth_per, .9), "encoded": len(ENC),
"supervised_ratio_mean": stats.mean(sup), "eval": results}, f, indent=2)
print(f"\n[9] wrote:\n {struct_path}\n {stats_path}")
print("done.")
Мы экспортируем каждую разобранную траекторию в виде структурированной записи JSONL, содержащей сообщения, схемы инструментов, вопросы и ответы. Мы также сохраняем JSON-отчёт, включающий размер корпуса, частотность инструментов, статистику траекторий, доли токенов под контролируемым обучением и доступные результаты оценки. Мы завершаем рабочий процесс подготовленными к повторному использованию артефактами датасета, аналитическими результатами и файлами модели, сохранёнными в настроенном выходном каталоге.
В заключение мы завершили практический конвейер анализа, преобразования, тонкой настройки и оценки сложных траекторий использования инструментов из датасета XYZ-Aquila-SFT. Мы сохранили исходную структуру диалогов, применили поуровневое обучение только к ответам ассистента и эффективно адаптировали Qwen3-0.6B с помощью LoRA на совместимом с Colab графическом процессоре. Мы также сравнили эффективность вызова инструментов до и после обучения с помощью оценки с принудительной подачей правильных ответов и экспортировали пригодные для повторного использования структурированные записи, адаптеры модели и аналитическую статистику. Этот рабочий процесс создаёт прочную основу для масштабирования контролируемой тонкой настройки с поддержкой инструментов, тестирования альтернативных политик длины последовательности и обучения более способных агентных языковых моделей.
Посмотрите ПОЛНЫЙ КОД здесь. Также подписывайтесь на нас в Twitter и не забудьте присоединиться к нашему ML-сабреддиту с более чем 150 тыс. участников и подписаться на нашу рассылку. Постойте! Вы есть в Telegram? теперь вы также можете присоединиться к нам в Telegram.
Хотите сотрудничать с нами для продвижения вашего репозитория GitHub, страницы Hugging Face, выпуска продукта, вебинара и т. д.? Свяжитесь с нами
Переведено автоматически с английского. Оригинал статьи — по ссылке ниже.