""" 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()