Buckets:
| """SFT jsonl -> chat-template tokens + assistant-only labels -> packed memmaps. | |
| data/pack.py cannot be used for SFT. It reads `json.loads(line)["text"]`, and an | |
| SFT row has `messages`, so it raises KeyError on the first line. More | |
| importantly it produces no loss mask, and train.py then calls | |
| `model(input_ids=ids, labels=ids)` -- loss on every token. For pretraining that | |
| is correct. For SFT it trains the model to generate the user's turns and the | |
| tool results as eagerly as its own replies. | |
| So this does three things pack.py does not: | |
| 1. RENDERS each row through the model's own chat_template, rather than | |
| concatenating raw text. The template is what decides whether a row becomes | |
| <think>reasoning</think>answer or <think></think>answer, so it is also what | |
| decides whether this model ends up hybrid. | |
| 2. MASKS every token that is not inside an assistant turn. Assistant spans run | |
| from <|im_start|>assistant to the matching <|im_end|>; everything else -- | |
| system prompt, tool schemas, user turns, tool responses -- gets label -100. | |
| Measured on this corpus that is 40% of all tokens, so it is not a detail. | |
| 3. GATES on real tokenizer length. assemble.py can only estimate from | |
| characters, and even at the corrected 3.44/1.33 ratios 19.1% of the 64K file | |
| still came out over its own cap, one row reaching 157,783 tokens against a | |
| 131,072 position limit. An estimate cannot fix that; only the tokenizer can. | |
| Output per shard: {stem}.bin (uint32 ids), {stem}.mask (uint8, 1 = compute loss), | |
| {stem}.json (meta). Labels are rebuilt as `where(mask, ids, -100)` at load time, | |
| which costs one byte per token instead of four. | |
| python data/pack_sft.py --input data/final_data_for_sft_training/train.jsonl \ | |
| --seq-len 8192 --out data/packed_sft | |
| python data/pack_sft.py --input data/final_data_for_sft_training/train_long64k.jsonl \ | |
| --seq-len 65536 --out data/packed_sft_long | |
| """ | |
| from __future__ import annotations | |
| import argparse | |
| import json | |
| import os | |
| import time | |
| from pathlib import Path | |
| import numpy as np | |
| os.environ.setdefault("TOKENIZERS_PARALLELISM", "true") | |
| from transformers import AutoTokenizer # noqa: E402 | |
| DTYPE = np.uint32 # Qwen3 vocab 151,936 does not fit uint16 | |
| REPO = Path(__file__).resolve().parent.parent | |
| def assistant_spans(ids, im_start, im_end, asst_ids): | |
| """Byte mask over ids: 1 inside an assistant turn, 0 everywhere else. | |
| The template writes '<|im_start|>assistant\\n ... <|im_end|>'. The role name | |
| is matched explicitly rather than assuming assistant turns alternate, | |
| because tool conversations do not alternate -- a tool response is rendered | |
| as a user turn, so position parity is meaningless here. | |
| """ | |
| mask = np.zeros(len(ids), dtype=np.uint8) | |
| n = len(asst_ids) | |
| i = 0 | |
| while i < len(ids): | |
| if ids[i] == im_start and ids[i + 1:i + 1 + n] == asst_ids: | |
| j = i + 1 + n | |
| while j < len(ids) and ids[j] != im_end: | |
| j += 1 | |
| # start after the role header so '<|im_start|>assistant' itself is | |
| # not a prediction target; include <|im_end|> so the model learns to | |
| # stop. Without that the model never emits a stop token. | |
| mask[i + 1 + n:min(j + 1, len(ids))] = 1 | |
| i = j | |
| i += 1 | |
| return mask | |
| def main(): | |
| ap = argparse.ArgumentParser() | |
| ap.add_argument("--input", required=True) | |
| ap.add_argument("--out", default="data/packed_sft") | |
| ap.add_argument("--tokenizer", default=str(REPO / "model")) | |
| ap.add_argument("--seq-len", type=int, default=8192) | |
| ap.add_argument("--batch", type=int, default=256) | |
| ap.add_argument("--limit", type=int, default=0, help="stop after N rows (smoke test)") | |
| args = ap.parse_args() | |
| inp = Path(args.input) | |
| outdir = Path(args.out) | |
| outdir.mkdir(parents=True, exist_ok=True) | |
| stem = inp.stem | |
| bin_path, mask_path = outdir / f"{stem}.bin", outdir / f"{stem}.mask" | |
| meta_path = outdir / f"{stem}.json" | |
| tok = AutoTokenizer.from_pretrained(args.tokenizer) | |
| eos = tok.eos_token_id | |
| if eos is None: | |
| raise SystemExit("tokenizer has no eos_token_id") | |
| im_start = tok.convert_tokens_to_ids("<|im_start|>") | |
| im_end = tok.convert_tokens_to_ids("<|im_end|>") | |
| asst_ids = tok("assistant", add_special_tokens=False)["input_ids"] | |
| print(f"{stem}: seq_len {args.seq_len:,} tokenizer {Path(args.tokenizer).name}") | |
| print(f" <|im_start|>={im_start} <|im_end|>={im_end} assistant={asst_ids}\n", | |
| flush=True) | |
| buf_ids: list[int] = [] | |
| buf_mask: list[int] = [] | |
| n_rows = n_seqs = n_tok = n_asst = 0 | |
| n_over = n_fail = n_read = 0 | |
| fail_ex: list[str] = [] | |
| t0 = time.time() | |
| def flush(fi, fm): | |
| nonlocal n_seqs, buf_ids, buf_mask | |
| usable = (len(buf_ids) // args.seq_len) * args.seq_len | |
| if usable == 0: | |
| return | |
| fi.write(np.asarray(buf_ids[:usable], dtype=DTYPE).tobytes()) | |
| fm.write(np.asarray(buf_mask[:usable], dtype=np.uint8).tobytes()) | |
| n_seqs += usable // args.seq_len | |
| buf_ids = buf_ids[usable:] | |
| buf_mask = buf_mask[usable:] | |
| with open(bin_path, "wb") as fi, open(mask_path, "wb") as fm: | |
| pend: list[str] = [] | |
| def drain(): | |
| nonlocal n_rows, n_tok, n_asst, n_over | |
| if not pend: | |
| return | |
| for ids in tok(pend, add_special_tokens=False)["input_ids"]: | |
| # Hard gate on REAL length. Dropping beats truncating: cutting a | |
| # row at seq_len removes the final assistant turn, which for the | |
| # long-context rows is the entire answer -- the model would be | |
| # trained on a document with no reply attached. | |
| if len(ids) > args.seq_len: | |
| n_over += 1 | |
| continue | |
| m = assistant_spans(ids, im_start, im_end, asst_ids) | |
| buf_ids.extend(ids) | |
| buf_mask.extend(m.tolist()) | |
| buf_ids.append(eos) | |
| buf_mask.append(1) # learn to stop at the row boundary | |
| n_rows += 1 | |
| n_tok += len(ids) + 1 | |
| n_asst += int(m.sum()) + 1 | |
| pend.clear() | |
| flush(fi, fm) | |
| with open(inp, encoding="utf-8") as f: | |
| for line in f: | |
| if args.limit and n_read >= args.limit: | |
| break | |
| n_read += 1 | |
| try: | |
| r = json.loads(line) | |
| s = tok.apply_chat_template( | |
| r["messages"], tools=r.get("tools") or None, | |
| tokenize=False, add_generation_prompt=False) | |
| except Exception as e: # noqa: BLE001 | |
| n_fail += 1 | |
| if len(fail_ex) < 5: | |
| fail_ex.append(f"line {n_read}: {type(e).__name__}: {str(e)[:70]}") | |
| continue | |
| pend.append(s) | |
| if len(pend) >= args.batch: | |
| drain() | |
| if n_rows and n_rows % 50_000 < args.batch: | |
| el = time.time() - t0 | |
| print(f" {n_rows:,} rows {n_tok/1e6:.0f}M tok " | |
| f"{n_tok/max(el,1)/1000:.0f}k tok/s", flush=True) | |
| drain() | |
| packed = n_seqs * args.seq_len | |
| meta = { | |
| "source": stem, | |
| "bin": str(bin_path).replace("\\", "/"), | |
| "mask": str(mask_path).replace("\\", "/"), | |
| "seq_len": args.seq_len, | |
| "n_sequences": n_seqs, | |
| "n_tokens": packed, | |
| "n_rows": n_rows, | |
| "n_tokens_seen": n_tok, | |
| "n_assistant_tokens": n_asst, | |
| "assistant_fraction": round(n_asst / max(n_tok, 1), 4), | |
| "rows_over_seq_len_dropped": n_over, | |
| "rows_failed_render": n_fail, | |
| "dtype": "uint32", | |
| "mask_dtype": "uint8", | |
| "tokenizer": args.tokenizer, | |
| } | |
| meta_path.write_text(json.dumps(meta, indent=2), encoding="utf-8") | |
| # Merge into the manifest train/sft_data.py loads. Re-packing one shard | |
| # updates its entry in place rather than appending a duplicate, so the | |
| # command can be re-run per input file in any order. | |
| man_path = outdir / "manifest.json" | |
| man = {"seq_len": args.seq_len, "sources": []} | |
| if man_path.exists(): | |
| man = json.loads(man_path.read_text(encoding="utf-8")) | |
| if man.get("seq_len") != args.seq_len: | |
| raise SystemExit( | |
| f"{man_path} holds seq_len {man.get('seq_len')} but this shard is " | |
| f"{args.seq_len}. train/sft_train.py requires one seq_len per " | |
| f"manifest -- pack different lengths into different --out dirs.") | |
| man["sources"] = [s for s in man["sources"] if s["name"] != stem] | |
| man["sources"].append({ | |
| "name": stem, | |
| "bin": meta["bin"], | |
| "mask": meta["mask"], | |
| "n_sequences": n_seqs, | |
| "n_tokens": packed, | |
| }) | |
| man["sources"].sort(key=lambda s: s["name"]) | |
| man_path.write_text(json.dumps(man, indent=2), encoding="utf-8") | |
| el = time.time() - t0 | |
| print(f"\n rows packed {n_rows:>12,}") | |
| print(f" render failures {n_fail:>12,}" | |
| f" {'clean' if not n_fail else 'INVESTIGATE'}") | |
| for x in fail_ex: | |
| print(f" {x}") | |
| print(f" dropped over seq_len {n_over:>12,} " | |
| f"({n_over/max(n_rows+n_over,1)*100:.2f}%)") | |
| print(f" tokens packed {packed:>12,} ({packed/1e9:.3f}B)") | |
| print(f" assistant tokens {n_asst:>12,} " | |
| f"({meta['assistant_fraction']*100:.1f}% carry loss)") | |
| print(f" sequences {n_seqs:>12,} x {args.seq_len:,}") | |
| print(f" elapsed {el/60:>12.1f} min " | |
| f"({n_tok/max(el,1)/1000:.0f}k tok/s)") | |
| print(f" -> {bin_path}\n -> {mask_path}\n -> {meta_path}") | |
| print(f" -> {man_path} ({len(man['sources'])} shard(s) registered)") | |
| if __name__ == "__main__": | |
| main() | |
Xet Storage Details
- Size:
- 10 kB
- Xet hash:
- 266631fa54708f4e3d18a3bbe212dd8c572b8dbd4d01681540dde2e9a834d4ac
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.