| """Datasets, collation (random mask ratio = masked-diffusion training), splits, negatives.""" |
| from __future__ import annotations |
|
|
| import json |
| import random |
|
|
| import torch |
|
|
| from .features import build_pass_a, build_pass_b, gold_plan |
| from .schema import Example, example_from_dict, example_to_dict |
|
|
|
|
| def load_jsonl(path) -> list[Example]: |
| with open(path, encoding="utf-8") as f: |
| return [example_from_dict(json.loads(line)) for line in f if line.strip()] |
|
|
|
|
| def save_jsonl(path, examples: list[Example]): |
| with open(path, "w", encoding="utf-8") as f: |
| for ex in examples: |
| f.write(json.dumps(example_to_dict(ex), ensure_ascii=False) + "\n") |
|
|
|
|
| def featurize(tok, ex: Example, cached: bool = False, markers: bool = False, v3: bool = False) -> dict | None: |
| if v3: |
| from .v3feat import build_pass_a_v3, build_pass_b_v3, gold_plan_v3 |
| a = build_pass_a_v3(tok, ex) |
| if a is None: |
| return None |
| b = None |
| if ex.calls and len(ex.calls) <= len(a["call_pos"]): |
| b = build_pass_b_v3(tok, ex, a, gold_plan_v3(tok, ex, a), with_targets=True, keep_context=cached, |
| markers=markers) |
| return {"a": a, "b": b} |
| a = build_pass_a(tok, ex) |
| if a is None: |
| return None |
| b = None |
| if ex.calls and len(ex.calls) <= len(a["call_pos"]): |
| b = build_pass_b(tok, ex, a, gold_plan(a, ex), with_targets=True, keep_context=cached, markers=markers) |
| return {"a": a, "b": b} |
|
|
|
|
| class FeatureDataset(torch.utils.data.Dataset): |
| def __init__(self, tok, examples: list[Example], cached: bool = False, markers: bool = False, v3: bool = False): |
| self.tok, self.examples, self.cached, self.markers, self.v3 = tok, examples, cached, markers, v3 |
|
|
| def __len__(self): |
| return len(self.examples) |
|
|
| def __getitem__(self, i): |
| return featurize(self.tok, self.examples[i], self.cached, self.markers, self.v3) |
|
|
| def length(self, i) -> int: |
| """Cheap length proxy for bucketing (chars/3).""" |
| ex = self.examples[i] |
| n = len(ex.query) + sum(len(t.name) + len(t.description) + 20 * len(t.params) for t in ex.tools) |
| return n // 3 + 40 |
|
|
|
|
| def _pad(seqs, value): |
| L = max(len(s) for s in seqs) |
| ids = torch.full((len(seqs), L), value, dtype=torch.long) |
| attn = torch.zeros((len(seqs), L), dtype=torch.long) |
| for i, s in enumerate(seqs): |
| ids[i, : len(s)] = torch.tensor(s) |
| attn[i, : len(s)] = 1 |
| return ids, attn |
|
|
|
|
| def _mask_slots(b, rng, v3_mask: str = "clean"): |
| """Masked-diffusion corruption: random ratio r ~ U(0,1], mask that share of slot tokens, reveal the rest. |
| v3_mask="legacy" reproduces the original v3 run (s10): every anchor always masked, gate trained on all anchors.""" |
| ids = list(b["input_ids"]) |
| flat = [(p, t) for pos, tg in zip(b["slot_pos"], b["slot_targets"]) for p, t in zip(pos, tg)] |
| if not flat: |
| return ids, [] |
| n_mask = max(1, round(rng.uniform(0.0, 1.0) * len(flat))) |
| masked = set(rng.sample(range(len(flat)), n_mask)) |
| always = set() |
| if b.get("anchors") and v3_mask == "legacy": |
| always = set(b["anchors"]) |
| b.pop("_clean", None) |
| elif b.get("anchors"): |
| |
| |
| |
| clean = [pos for pos in b["slot_pos"] if rng.random() < 0.5] |
| always = {p for pos in clean for p in pos} |
| b["_clean"] = [pos[0] for pos in clean] |
| targets = [] |
| for idx, (p, t) in enumerate(flat): |
| if idx in masked or p in always: |
| targets.append((p, t)) |
| else: |
| ids[p] = t |
| return ids, targets |
|
|
|
|
| def collate(items, tok, rng: random.Random | None = None, cached: bool = False, v3_mask: str = "clean"): |
| rng = rng or random |
| items = [it for it in items if it is not None] |
| if not items: |
| return None |
| if cached: |
| return _collate_cached(items, tok, rng, v3_mask) |
| A = [it["a"] for it in items] |
| a_ids, a_attn = _pad([a["input_ids"] for a in A], tok.pad_token_id) |
| J = max(max(len(a["t_pos"]) for a in A), 1) |
| t_pos = torch.full((len(A), J), -1, dtype=torch.long) |
| for i, a in enumerate(A): |
| t_pos[i, : len(a["t_pos"])] = torch.tensor(a["t_pos"], dtype=torch.long) |
| dec = torch.tensor([[J if (l == len(a["t_pos"])) else l for l in a["dec_labels"]] for a in A]) |
| len_idx, len_labels = [], [] |
| for i, a in enumerate(A): |
| for (ci, j, k, b) in a["len_items"]: |
| len_idx.append((i, a["call_pos"][ci], a["arg_pos"][(j, k)])) |
| len_labels.append(b) |
| batch = {"a_ids": a_ids, "a_attn": a_attn, "t_pos": t_pos, "t_mask": t_pos >= 0, |
| "call_pos": torch.tensor([a["call_pos"] for a in A]), "dec_labels": dec, |
| "len_idx": torch.tensor(len_idx, dtype=torch.long).reshape(-1, 3), |
| "len_labels": torch.tensor(len_labels, dtype=torch.long)} |
|
|
| |
| B = [it["b"] for it in items if it["b"] is not None and (it["b"]["slot_pos"] or it["b"].get("opt_slots"))] |
| if B: |
| seqs, sb, sp, sl = [], [], [], [] |
| for bi, b in enumerate(B): |
| ids, targets = _mask_slots(b, rng, v3_mask) |
| for p, t in targets: |
| sb.append(bi); sp.append(p); sl.append(t) |
| seqs.append(ids) |
| b_ids, b_attn = _pad(seqs, tok.pad_token_id) |
| batch.update({"b_ids": b_ids, "b_attn": b_attn, "slot_b": torch.tensor(sb), "slot_p": torch.tensor(sp), |
| "slot_labels": torch.tensor(sl)}) |
| else: |
| batch["b_ids"] = None |
| if "q_pos" in A[0]: |
| from .v3feat import v3_extras |
| batch.update(v3_extras(A, list(enumerate(B)))) |
| return batch |
|
|
|
|
| def _collate_cached(items, tok, rng, v3_mask="clean"): |
| """One sequence per example: pass-A context + slot block, read with the block-causal mask.""" |
| A, seqs, ctx, sb, sp, sl = [], [], [], [], [], [] |
| for i, it in enumerate(items): |
| a, b = it["a"], it["b"] |
| A.append(a) |
| ctx.append(len(a["input_ids"])) |
| if b is not None and (b["slot_pos"] or b.get("opt_slots")): |
| ids, targets = _mask_slots(b, rng, v3_mask) |
| for p, t in targets: |
| sb.append(i); sp.append(p); sl.append(t) |
| seqs.append(ids) |
| else: |
| seqs.append(a["input_ids"]) |
| ids, attn = _pad(seqs, tok.pad_token_id) |
| J = max(max(len(a["t_pos"]) for a in A), 1) |
| t_pos = torch.full((len(A), J), -1, dtype=torch.long) |
| for i, a in enumerate(A): |
| t_pos[i, : len(a["t_pos"])] = torch.tensor(a["t_pos"], dtype=torch.long) |
| dec = torch.tensor([[J if (l == len(a["t_pos"])) else l for l in a["dec_labels"]] for a in A]) |
| len_idx, len_labels = [], [] |
| for i, a in enumerate(A): |
| for (ci, j, k, b) in a["len_items"]: |
| len_idx.append((i, a["call_pos"][ci], a["arg_pos"][(j, k)])) |
| len_labels.append(b) |
| out = {"a_ids": ids, "a_attn": attn, "t_pos": t_pos, "t_mask": t_pos >= 0, |
| "call_pos": torch.tensor([a["call_pos"] for a in A]), "dec_labels": dec, |
| "len_idx": torch.tensor(len_idx, dtype=torch.long).reshape(-1, 3), |
| "len_labels": torch.tensor(len_labels, dtype=torch.long), |
| "ctx_len": torch.tensor(ctx), "b_ids": None, "slot_b": torch.tensor(sb, dtype=torch.long), |
| "slot_p": torch.tensor(sp, dtype=torch.long), "slot_labels": torch.tensor(sl, dtype=torch.long)} |
| if "q_pos" in A[0]: |
| from .v3feat import v3_extras |
| out.update(v3_extras(A, [(i, it["b"]) for i, it in enumerate(items) |
| if it["b"] is not None and (it["b"]["slot_pos"] or it["b"].get("opt_slots"))])) |
| return out |
|
|
|
|
| def token_batches(lengths: list[int], tokens_per_batch: int, rng: random.Random, max_bs: int = 256): |
| """Length-bucketed batches (sorted by length, chunked by token budget, shuffled).""" |
| order = sorted(range(len(lengths)), key=lambda i: lengths[i]) |
| batches, cur, cur_max = [], [], 0 |
| for i in order: |
| m = max(cur_max, lengths[i]) |
| if cur and (m * (len(cur) + 1) > tokens_per_batch or len(cur) >= max_bs): |
| batches.append(cur); cur, m = [], lengths[i] |
| cur.append(i); cur_max = m |
| if cur: |
| batches.append(cur) |
| rng.shuffle(batches) |
| return batches |
|
|
|
|
| def hard_negatives(examples: list[Example], n: int, rng: random.Random) -> list[Example]: |
| """Same query, same domain: drop every tool the gold answer uses, keep its related neighbours -> no call.""" |
| out = [] |
| pool = [e for e in examples if e.calls and len(e.tools) > len({c.name for c in e.calls})] |
| rng.shuffle(pool) |
| for e in pool[:n]: |
| needed = {c.name for c in e.calls} |
| tools = [t for t in e.tools if t.name not in needed] |
| out.append(Example(e.query, tools, [], id=f"hardneg_{len(out)}", source="xlam_hardneg", gold=[])) |
| return out |
|
|
|
|
| def irrelevance_negatives(examples: list[Example], n: int, rng: random.Random) -> list[Example]: |
| """Pair real queries with tool lists from other examples that don't contain any needed tool.""" |
| out = [] |
| pool = [e for e in examples if e.tools] |
| while len(out) < n: |
| q, other = rng.choice(pool), rng.choice(pool) |
| needed = {c.name for c in q.calls} |
| tools = [t for t in other.tools if t.name not in needed] |
| if not tools or other is q: |
| continue |
| out.append(Example(q.query, tools, [], id=f"neg_{len(out)}", source="xlam_neg", gold=[])) |
| return out |
|
|
