| """Convert text or chat JSONL into streamable uint32 tokens and uint8 training masks. |
| Chat JSONL: {"messages": [{"role":"user", "content":"..."}, {"role":"assistant", "content":"..."}]} |
| Only assistant answers contribute to SFT loss. Input messages must be trusted training data. |
| """ |
| import argparse |
| import json |
| from pathlib import Path |
| from array import array |
| from tokenizers import Tokenizer |
|
|
|
|
| def sources(paths, fmt): |
| suffix = ".txt" if fmt == "text" else ".jsonl" |
| for src in paths: |
| p = Path(src) |
| if p.is_dir(): |
| yield from sorted(p.rglob("*")) if False else (x for x in sorted(p.rglob(f"*{suffix}")) if x.is_file()) |
| elif p.is_file(): |
| yield p |
| else: |
| raise FileNotFoundError(src) |
|
|
|
|
| def main(): |
| p = argparse.ArgumentParser() |
| p.add_argument("--input", nargs="+", required=True) |
| p.add_argument("--tokenizer", default="tokenizer.json") |
| p.add_argument("--output-dir", required=True) |
| p.add_argument("--format", choices=["text", "chat"], default="text") |
| args = p.parse_args() |
| tok = Tokenizer.from_file(args.tokenizer) |
| eos = tok.token_to_id("<|eos|>") |
| end = tok.token_to_id("<|end|>") |
| if eos is None or end is None: |
| raise ValueError("Tokenizer special tokens are missing") |
| out = Path(args.output_dir) |
| out.mkdir(parents=True, exist_ok=True) |
| count = 0 |
| with (out / "train.bin").open("wb") as ids_file, (out / "mask.bin").open("wb") as mask_file: |
| def emit(ids, mask): |
| nonlocal count |
| if not ids: |
| return |
| array("I", ids).tofile(ids_file) |
| mask_file.write(bytes([mask]) * len(ids)) |
| count += len(ids) |
|
|
| for path in sources(args.input, args.format): |
| with path.open("r", encoding="utf-8") as f: |
| if args.format == "text": |
| |
| while True: |
| chunk = f.read(256_000) |
| if not chunk: |
| break |
| emit(tok.encode(chunk, add_special_tokens=False).ids, 1) |
| emit([eos], 1) |
| else: |
| for lineno, line in enumerate(f, 1): |
| if not line.strip(): |
| continue |
| try: |
| record = json.loads(line) |
| messages = record["messages"] |
| for m in messages: |
| role = m["role"] |
| if role not in ("system", "user", "assistant"): |
| raise ValueError(f"Unsupported role {role}") |
| rid = tok.token_to_id(f"<|{role}|>") |
| emit([rid], 0) |
| emit(tok.encode(str(m["content"]), add_special_tokens=False).ids, |
| 1 if role == "assistant" else 0) |
| emit([end], 1 if role == "assistant" else 0) |
| emit([eos], 0) |
| except (KeyError, ValueError, TypeError) as exc: |
| raise ValueError(f"{path}:{lineno}: {exc}") from exc |
| if count < 3: |
| raise ValueError("Too few tokens. Supply more training data.") |
| (out / "metadata.json").write_text(json.dumps({"tokens": count, "format": args.format}, indent=2)) |
| print(f"Wrote {count:,} tokens to {out} (train.bin and mask.bin)") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
