File size: 8,429 Bytes
4d5ab1c | 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 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 | """
Preprocess local code-message JSONL datasets for llm.c GPT-2 training.
Expected inputs:
- text-only JSONL: {"text": "..."}
- messages JSONL: {"messages": [{"role": "...", "content": "..."}, ...]}
The output is llm.c's GPT-2 data format:
- 256 int32 header values (1024 bytes)
- followed by uint16 GPT-2 token ids
"""
import argparse
import glob
import hashlib
import json
import multiprocessing as mp
import os
from pathlib import Path
import numpy as np
import tiktoken
from tqdm import tqdm
HEADER_SIZE = 256
GPT2_DATA_MAGIC = 20240520
GPT2_DATA_VERSION = 1
_ENC = None
_EOT = None
def init_worker():
global _ENC, _EOT
_ENC = tiktoken.get_encoding("gpt2")
_EOT = _ENC._special_tokens["<|endoftext|>"]
def tokenize_text(text):
ids = _ENC.encode_ordinary(text)
tokens = np.empty(len(ids) + 1, dtype=np.uint16)
tokens[0] = _EOT
tokens[1:] = ids
return tokens
def tokenize_record(record):
doc_idx, text = record
return doc_idx, tokenize_text(text)
def write_datafile_np(filename, tokens):
assert tokens.dtype == np.uint16
assert len(tokens) < 2**31, "token count too large for one shard"
header = np.zeros(HEADER_SIZE, dtype=np.int32)
header[0] = GPT2_DATA_MAGIC
header[1] = GPT2_DATA_VERSION
header[2] = len(tokens)
num_bytes = HEADER_SIZE * 4 + len(tokens) * tokens.itemsize
print(f"writing {len(tokens):,} tokens to {filename} ({num_bytes:,} bytes)")
with open(filename, "wb") as f:
f.write(header.tobytes())
f.write(tokens.tobytes())
class ShardWriter:
def __init__(self, output_dir, prefix, shard_size):
self.output_dir = Path(output_dir)
self.prefix = prefix
self.shard_size = shard_size
self.shard_index = 0
self.token_count = 0
self.total_tokens = 0
self.buffer = None
def _ensure_buffer(self):
if self.buffer is None:
self.buffer = np.empty(self.shard_size, dtype=np.uint16)
def _filename(self):
return self.output_dir / f"{self.prefix}_{self.shard_index:06d}.bin"
def write(self, tokens):
if len(tokens) == 0:
return
self._ensure_buffer()
offset = 0
while offset < len(tokens):
space = self.shard_size - self.token_count
take = min(space, len(tokens) - offset)
self.buffer[self.token_count:self.token_count + take] = tokens[offset:offset + take]
self.token_count += take
self.total_tokens += take
offset += take
if self.token_count == self.shard_size:
write_datafile_np(self._filename(), self.buffer)
self.shard_index += 1
self.token_count = 0
def close(self):
if self.buffer is not None and self.token_count > 0:
write_datafile_np(self._filename(), self.buffer[:self.token_count].copy())
self.shard_index += 1
self.token_count = 0
def serialize_messages(messages):
parts = []
for message in messages:
role = str(message.get("role", "unknown"))
content = str(message.get("content", ""))
if content:
parts.append(f"<|{role}|>\n{content}")
return "\n".join(parts)
def iter_texts(path, input_format, text_key, messages_key, limit_docs):
skipped = 0
yielded = 0
with open(path, "r", encoding="utf-8") as f:
for line_idx, line in enumerate(f):
if limit_docs is not None and yielded >= limit_docs:
break
try:
obj = json.loads(line)
except json.JSONDecodeError:
skipped += 1
continue
if input_format == "text":
text = obj.get(text_key, "")
elif input_format == "messages":
text = serialize_messages(obj.get(messages_key, []))
else:
if text_key in obj:
text = obj.get(text_key, "")
elif messages_key in obj:
text = serialize_messages(obj.get(messages_key, []))
else:
text = ""
if not isinstance(text, str):
text = str(text)
if not text.strip():
skipped += 1
continue
yield line_idx, text
yielded += 1
if skipped:
print(f"skipped {skipped:,} empty or invalid records")
def goes_to_val(doc_index, val_fraction, seed):
if val_fraction <= 0:
return False
key = f"{seed}:{doc_index}".encode("utf-8")
digest = hashlib.blake2b(key, digest_size=8).digest()
value = int.from_bytes(digest, "little") / 2**64
return value < val_fraction
def ensure_no_existing_bins(output_dir, dataset_name, overwrite):
patterns = [
os.path.join(output_dir, f"{dataset_name}_train_*.bin"),
os.path.join(output_dir, f"{dataset_name}_val_*.bin"),
]
existing = [path for pattern in patterns for path in glob.glob(pattern)]
if existing and not overwrite:
raise SystemExit(
f"Refusing to overwrite {len(existing)} existing .bin files in {output_dir}. "
"Pass --overwrite to replace them."
)
for path in existing:
os.remove(path)
def main():
parser = argparse.ArgumentParser(description="Preprocess code-message JSONL for llm.c GPT-2 training")
parser.add_argument("--input", required=True, help="Input JSONL file")
parser.add_argument("--output_dir", default=None, help="Output directory for .bin shards")
parser.add_argument("--dataset_name", default="code_messages", help="Prefix for output shard files")
parser.add_argument("--format", choices=["auto", "text", "messages"], default="auto", help="Input JSONL schema")
parser.add_argument("--text_key", default="text", help="Text field name for text JSONL")
parser.add_argument("--messages_key", default="messages", help="Messages field name for chat JSONL")
parser.add_argument("--shard_size", type=int, default=100_000_000, help="Tokens per output shard")
parser.add_argument("--val_fraction", type=float, default=0.001, help="Doc fraction to reserve for validation")
parser.add_argument("--seed", type=int, default=1337, help="Seed for deterministic validation split")
parser.add_argument("--workers", type=int, default=min(16, max(1, (os.cpu_count() or 2) - 2)))
parser.add_argument("--chunksize", type=int, default=16, help="Multiprocessing chunksize")
parser.add_argument("--limit_docs", type=int, default=None, help="Only preprocess this many docs, for smoke tests")
parser.add_argument("--total_docs", type=int, default=None, help="Optional tqdm total")
parser.add_argument("--overwrite", action="store_true", help="Delete existing output shards first")
args = parser.parse_args()
if not (0.0 <= args.val_fraction < 1.0):
raise SystemExit("--val_fraction must be in [0, 1)")
if args.shard_size <= 0:
raise SystemExit("--shard_size must be positive")
script_dir = Path(__file__).resolve().parent
output_dir = Path(args.output_dir) if args.output_dir else script_dir / args.dataset_name
output_dir.mkdir(parents=True, exist_ok=True)
ensure_no_existing_bins(str(output_dir), args.dataset_name, args.overwrite)
train_writer = ShardWriter(output_dir, f"{args.dataset_name}_train", args.shard_size)
val_writer = ShardWriter(output_dir, f"{args.dataset_name}_val", args.shard_size)
total = args.total_docs
if args.limit_docs is not None:
total = args.limit_docs if total is None else min(total, args.limit_docs)
with mp.Pool(args.workers, initializer=init_worker) as pool:
records = iter_texts(args.input, args.format, args.text_key, args.messages_key, args.limit_docs)
token_iter = pool.imap(tokenize_record, records, chunksize=args.chunksize)
for doc_idx, tokens in tqdm(token_iter, total=total, unit="docs"):
if goes_to_val(doc_idx, args.val_fraction, args.seed):
val_writer.write(tokens)
else:
train_writer.write(tokens)
train_writer.close()
val_writer.close()
print(f"train tokens: {train_writer.total_tokens:,}")
print(f"val tokens: {val_writer.total_tokens:,}")
print(f"wrote shards under: {output_dir}")
if __name__ == "__main__":
main()
|