File size: 3,623 Bytes
7da1a03 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 | """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":
# Read in chunks while preserving valid UTF-8 decoding.
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()
|