mport os import json import numpy as np from tokenizers import Tokenizer from tqdm import tqdm import multiprocessing as mp from functools import partial import time import pickle import psutil import signal import sys from pathlib import Path from collections import deque import gc import gzip

import os os.environ["TOKENIZERS_PARALLELISM"] = "false"

--- Configuration --- TOKENIZER_PATH = "tokenizer.json" RAW_JSONL_FILES_DIR = "./DATA_CHUNKS" TEXT_FIELD_NAME = "text" OUTPUT_BIN_DIR = "./data_tokenized" CHECKPOINT_DIR = "./checkpoints" SHARD_SIZE_TOKENS = 100 * 1000 * 1000 # 100 million tokens BIN_MAGIC_NUMBER = 20240520 BIN_VERSION = 1

--- Multiprocessing Configuration --- NUM_PROCESSES = 65 # Leave some cores for OS LINES_PER_TASK_CHUNK = 4096 POOL_DISPATCH_CHUNKSIZE = max(1, NUM_PROCESSES // 4) MAX_TASKS_PER_WORKER = 1000 # Recycle workers after this many tasks

--- Memory Management --- MAX_BUFFER_SIZE_GB = 20 # Maximum size of token buffer in GB MEMORY_CHECK_INTERVAL = 10 # Check memory every N chunks (reduced from 100) MEMORY_USAGE_WARNING_THRESHOLD = 0.85 # Warn if using >85% of available memory GC_INTERVAL = 50 # Force garbage collection every N chunks

--- Checkpoint Configuration --- CHECKPOINT_INTERVAL = 50 # Save checkpoint every N chunks processed ENABLE_LINE_COUNTING = False # Set to False for very large files USE_COMPRESSED_CHECKPOINTS = True # Use gzip compression for checkpoints

--- Global variables for workers --- worker_tokenizer_instance = None worker_eot_token_id_global = None worker_text_field_name_global = None

--- Checkpoint State --- class CheckpointState: def init(self): self.completed_files = set() self.current_file = None self.current_file_position = 0 # Byte position in current file self.current_file_line = 0 # Line number in current file self.shard_index = 0 self.total_docs_processed = 0 self.all_tokens_buffer = deque() # Changed from list to deque self.start_time = time.time()

def save(self, checkpoint_path): """Save checkpoint to disk with optional compression""" os.makedirs(os.path.dirname(checkpoint_path), exist_ok=True)

# Convert deque to list for pickling (maintains compatibility)
temp_buffer = list(self.all_tokens_buffer)
self.all_tokens_buffer = temp_buffer

try:
    if USE_COMPRESSED_CHECKPOINTS:
        with gzip.open(checkpoint_path + '.gz', 'wb', compresslevel=3) as f:
            pickle.dump(self, f)
    else:
        with open(checkpoint_path, 'wb') as f:
            pickle.dump(self, f)
finally:
    # Convert back to deque
    self.all_tokens_buffer = deque(self.all_tokens_buffer)

# Also save a human-readable status
status_path = checkpoint_path.replace('.pkl', '_status.json')
with open(status_path, 'w') as f:
    json.dump({
        'completed_files': list(self.completed_files),
        'current_file': self.current_file,
        'current_file_line': self.current_file_line,
        'shard_index': self.shard_index,
        'total_docs_processed': self.total_docs_processed,
        'buffer_size': len(self.all_tokens_buffer),
        'buffer_size_mb': (len(self.all_tokens_buffer) * 2) / (1024**2),  # uint16 = 2 bytes
        'elapsed_time': time.time() - self.start_time
    }, f, indent=2)
@classmethod def load(cls, checkpoint_path): """Load checkpoint from disk""" if os.path.exists(checkpoint_path + '.gz'): with gzip.open(checkpoint_path + '.gz', 'rb') as f: state = pickle.load(f) # Ensure buffer is a deque if isinstance(state.all_tokens_buffer, list): state.all_tokens_buffer = deque(state.all_tokens_buffer) return state elif os.path.exists(checkpoint_path): with open(checkpoint_path, 'rb') as f: state = pickle.load(f) # Ensure buffer is a deque if isinstance(state.all_tokens_buffer, list): state.all_tokens_buffer = deque(state.all_tokens_buffer) return state return None --- Signal Handler for Graceful Shutdown --- shutdown_requested = False

def signal_handler(signum, frame): global shutdown_requested print(f"\nReceived signal {signum}. Initiating graceful shutdown...") shutdown_requested = True

signal.signal(signal.SIGINT, signal_handler) signal.signal(signal.SIGTERM, signal_handler)

--- Worker Functions --- def initialize_worker(tokenizer_path_param, eot_token_str_param, text_field_name_param): """Initializer for each worker process in the pool.""" global worker_tokenizer_instance, worker_eot_token_id_global, worker_text_field_name_global

worker_tokenizer_instance = Tokenizer.from_file(tokenizer_path_param) worker_eot_token_id_global = worker_tokenizer_instance.token_to_id(eot_token_str_param) worker_text_field_name_global = text_field_name_param if worker_eot_token_id_global is None: raise ValueError(f"Worker {os.getpid()}: EOT token '{eot_token_str_param}' not found.") def process_json_lines_batch(batch_data): """ Process a batch of JSON lines with line number tracking. batch_data is a tuple of (chunk_lines, starting_line_number) """ chunk_lines, start_line_num = batch_data

global worker_tokenizer_instance, worker_eot_token_id_global, worker_text_field_name_global if worker_tokenizer_instance is None: print(f"CRITICAL ERROR in worker {os.getpid()}: Tokenizer not initialized!") return [], start_line_num + len(chunk_lines)

processed_docs_tokens = [] texts_for_batch_encoding = []

for json_line_str in chunk_lines: json_line_str = json_line_str.strip() if not json_line_str: continue try: json_doc = json.loads(json_line_str) document_text = json_doc.get(worker_text_field_name_global) if document_text and isinstance(document_text, str) and document_text.strip(): texts_for_batch_encoding.append(document_text) except json.JSONDecodeError: continue except Exception as e: print(f"Worker {os.getpid()} error processing line: {e}") continue

if not texts_for_batch_encoding: return [], start_line_num + len(chunk_lines)

try: encoded_batch = worker_tokenizer_instance.encode_batch(texts_for_batch_encoding, add_special_tokens=False) for encoded_doc in encoded_batch: processed_docs_tokens.append([worker_eot_token_id_global] + encoded_doc.ids) except Exception as e: print(f"ERROR in worker {os.getpid()} during encode_batch: {e}") return [], start_line_num + len(chunk_lines)

return processed_docs_tokens, start_line_num + len(chunk_lines) def write_datafile(filename, tokens_np_uint16): """Saves token data as a .bin file.""" num_tokens = len(tokens_np_uint16) if num_tokens == 0: return False

header = np.zeros(256, dtype=np.int32) header[0] = BIN_MAGIC_NUMBER header[1] = BIN_VERSION header[2] = num_tokens

with open(filename, "wb") as f: f.write(header.tobytes()) f.write(tokens_np_uint16.tobytes()) return True def extract_tokens_from_deque(token_deque, count): """Efficiently extract a fixed number of tokens from deque""" tokens = [] for _ in range(min(count, len(token_deque))): tokens.append(token_deque.popleft()) return tokens

def write_all_complete_shards(all_tokens_buffer, shard_index, output_dir): """Write all complete shards from buffer""" shards_written = 0

while len(all_tokens_buffer) >= SHARD_SIZE_TOKENS: # Extract tokens efficiently tokens_to_write = extract_tokens_from_deque(all_tokens_buffer, SHARD_SIZE_TOKENS)

tokens_np = np.array(tokens_to_write, dtype=np.uint16)
split = "val" if shard_index == 0 else "train"
output_filename = os.path.join(output_dir, f"mydata_{split}_{shard_index:06d}.bin")

if write_datafile(output_filename, tokens_np):
    shards_written += 1
    shard_index += 1

# Explicitly delete large objects
del tokens_to_write
del tokens_np
return shard_index, shards_written def read_lines_in_chunks_with_resume(file_path, chunk_size, start_line=0): """Generator to read file from a specific line number.""" lines_chunk = [] current_line = 0

try: with open(file_path, "r", encoding="utf-8") as f: # Skip to start line if resuming while current_line < start_line: f.readline() current_line += 1

    for line in f:
        lines_chunk.append(line)
        current_line += 1
        
        if len(lines_chunk) == chunk_size:
            yield (lines_chunk, current_line - chunk_size)
            lines_chunk = []
            
    if lines_chunk:
        yield (lines_chunk, current_line - len(lines_chunk))
except Exception as e: print(f"Error reading file {file_path}: {e}") def estimate_lines_in_file(file_path, sample_size=10000): """Estimate number of lines based on file size and sample.""" file_size = os.path.getsize(file_path) if file_size == 0: return 0

bytes_read = 0 lines_read = 0

with open(file_path, 'rb') as f: while lines_read < sample_size and bytes_read < file_size: line = f.readline() if not line: break bytes_read += len(line) lines_read += 1

if lines_read == 0: return 0

avg_line_size = bytes_read / lines_read estimated_lines = int(file_size / avg_line_size) return estimated_lines def get_memory_usage(): """Get current memory usage as a fraction of available memory.""" process = psutil.Process() mem_info = process.memory_info() available_memory = psutil.virtual_memory().available total_memory = psutil.virtual_memory().total used_memory = total_memory - available_memory return used_memory / total_memory, mem_info.rss / (1024**3) # fraction, GB

def recreate_pool(old_pool, num_processes, init_func, init_args): """Safely recreate the multiprocessing pool""" if old_pool: old_pool.close() old_pool.join() return mp.Pool( processes=num_processes, initializer=init_func, initargs=init_args, maxtasksperchild=MAX_TASKS_PER_WORKER )

--- Main Processing Logic --- if name == "main": try: mp.set_start_method("fork", force=True) print("Using 'fork' start method for multiprocessing.") except RuntimeError: print("Fork not available, using default start method.")

print(f"Main process ID: {os.getpid()}") print(f"Using {NUM_PROCESSES} worker processes.") print(f"Memory limit for buffer: {MAX_BUFFER_SIZE_GB} GB") print(f"Worker recycling after {MAX_TASKS_PER_WORKER} tasks")

Load or create checkpoint
os.makedirs(CHECKPOINT_DIR, exist_ok=True) checkpoint_path = os.path.join(CHECKPOINT_DIR, "tokenization_checkpoint.pkl")

checkpoint = CheckpointState.load(checkpoint_path) if checkpoint: print(f"Resuming from checkpoint:") print(f" - Completed files: {len(checkpoint.completed_files)}") print(f" - Current file: {checkpoint.current_file}") print(f" - Current line: {checkpoint.current_file_line}") print(f" - Shard index: {checkpoint.shard_index}") print(f" - Total docs: {checkpoint.total_docs_processed}") print(f" - Buffer size: {len(checkpoint.all_tokens_buffer)} tokens") print(f" - Buffer memory: {(len(checkpoint.all_tokens_buffer) * 2) / (1024**2):.2f} MB") else: checkpoint = CheckpointState() print("Starting fresh (no checkpoint found)")

Tokenizer setup
if not os.path.exists(TOKENIZER_PATH): raise FileNotFoundError(f"Tokenizer not found at {TOKENIZER_PATH}")

_temp_tokenizer = Tokenizer.from_file(TOKENIZER_PATH) EOT_TOKEN_STR_MAIN = "
" _eot_id_main_check = _temp_tokenizer.token_to_id(EOT_TOKEN_STR_MAIN) if _eot_id_main_check is None: raise ValueError(f"EOT token '{EOT_TOKEN_STR_MAIN}' not found.") if _temp_tokenizer.get_vocab_size() >= 2**16: raise ValueError(f"Tokenizer vocab size too large for uint16 storage.") print(f"Tokenizer loaded. Vocab size: {_temp_tokenizer.get_vocab_size()}") del _temp_tokenizer

os.makedirs(OUTPUT_BIN_DIR, exist_ok=True)

Get list of files to process
raw_files = sorted([ os.path.join(RAW_JSONL_FILES_DIR, f) for f in os.listdir(RAW_JSONL_FILES_DIR) if f.endswith(".jsonl") or f.endswith(".json") ])

if not raw_files: print(f"No files found in {RAW_JSONL_FILES_DIR}") exit()

print(f"Found {len(raw_files)} files to process.")

Filter out already completed files
files_to_process = [f for f in raw_files if f not in checkpoint.completed_files] if checkpoint.current_file and checkpoint.current_file in files_to_process: # Make sure current file is first files_to_process.remove(checkpoint.current_file) files_to_process.insert(0, checkpoint.current_file)

print(f"Remaining files to process: {len(files_to_process)}")

Processing variables from checkpoint
shard_index = checkpoint.shard_index all_tokens_buffer = checkpoint.all_tokens_buffer total_docs_processed = checkpoint.total_docs_processed

chunks_since_checkpoint = 0 chunks_since_memory_check = 0 chunks_since_gc = 0 total_chunks_processed = 0

pool_init_args = (TOKENIZER_PATH, EOT_TOKEN_STR_MAIN, TEXT_FIELD_NAME)

Create initial pool with worker recycling
pool = recreate_pool(None, NUM_PROCESSES, initialize_worker, pool_init_args)

try: for file_idx, raw_file_path in enumerate(files_to_process): if shutdown_requested: print("Shutdown requested, saving checkpoint...") break

    # Determine starting line for this file
    start_line = 0
    if raw_file_path == checkpoint.current_file:
        start_line = checkpoint.current_file_line
        print(f"\nResuming file ({file_idx+1}/{len(files_to_process)}): {raw_file_path} from line {start_line}")
    else:
        print(f"\nStarting file ({file_idx+1}/{len(files_to_process)}): {raw_file_path}")
        checkpoint.current_file = raw_file_path
        checkpoint.current_file_line = 0
    
    # Estimate lines for progress bar
    if ENABLE_LINE_COUNTING:
        try:
            with open(raw_file_path, 'r', encoding='utf-8') as f:
                for _ in range(start_line):
                    f.readline()
                num_lines_remaining = sum(1 for _ in f)
        except:
            num_lines_remaining = None
    else:
        # Use estimation for large files
        estimated_total = estimate_lines_in_file(raw_file_path)
        num_lines_remaining = max(0, estimated_total - start_line) if estimated_total else None
        if estimated_total:
            file_size_gb = os.path.getsize(raw_file_path) / (1024**3)
            print(f"  File size: {file_size_gb:.2f} GB, Estimated total lines: {estimated_total:,}")
    
    # Create generator for chunks
    line_chunk_generator = read_lines_in_chunks_with_resume(
        raw_file_path, LINES_PER_TASK_CHUNK, start_line
    )
    
    num_chunks_estimate = (num_lines_remaining // LINES_PER_TASK_CHUNK if num_lines_remaining else None)
    
    # Process chunks
    results_iterator = pool.imap_unordered(
        process_json_lines_batch,
        line_chunk_generator,
        chunksize=POOL_DISPATCH_CHUNKSIZE
    )
    
    progress_desc = f"Processing {os.path.basename(raw_file_path)}"
    progress_bar = tqdm(results_iterator, total=num_chunks_estimate, desc=progress_desc)
    
    for result in progress_bar:
        if shutdown_requested:
            break
            
        list_of_doc_tokens, last_line_processed = result
        checkpoint.current_file_line = last_line_processed
        
        if list_of_doc_tokens:
            for doc_tokens in list_of_doc_tokens:
                if doc_tokens:
                    all_tokens_buffer.extend(doc_tokens)
                    total_docs_processed += 1
        
        # Memory check (more frequent)
        chunks_since_memory_check += 1
        if chunks_since_memory_check >= MEMORY_CHECK_INTERVAL:
            mem_fraction, process_gb = get_memory_usage()
            buffer_tokens = len(all_tokens_buffer)
            buffer_size_est_gb = (buffer_tokens * 2) / (1024**3)  # uint16 = 2 bytes
            
            if mem_fraction > MEMORY_USAGE_WARNING_THRESHOLD:
                print(f"\nWARNING: High memory usage: {mem_fraction*100:.1f}%, Process RSS: {process_gb:.2f} GB")
                print(f"Buffer size: {buffer_tokens:,} tokens (~{buffer_size_est_gb:.2f} GB)")
            
            # Write ALL complete shards if buffer is large
            if buffer_size_est_gb > MAX_BUFFER_SIZE_GB * 0.8:  # Start writing at 80% of limit
                print(f"\nBuffer approaching limit, writing all complete shards...")
                new_shard_index, shards_written = write_all_complete_shards(
                    all_tokens_buffer, shard_index, OUTPUT_BIN_DIR
                )
                if shards_written > 0:
                    print(f"  Wrote {shards_written} shards ({shard_index} to {new_shard_index-1})")
                    shard_index = new_shard_index
            
            chunks_since_memory_check = 0
        
        # Regular shard writing
        if len(all_tokens_buffer) >= SHARD_SIZE_TOKENS:
            new_shard_index, shards_written = write_all_complete_shards(
                all_tokens_buffer, shard_index, OUTPUT_BIN_DIR
            )
            if shards_written > 0:
                shard_index = new_shard_index
                progress_bar.set_postfix({
                    'shard': shard_index-1,
                    'docs': f"{total_docs_processed:,}",
                    'buf': f"{len(all_tokens_buffer)/1e6:.1f}M"
                })
        
        # Garbage collection
        chunks_since_gc += 1
        if chunks_since_gc >= GC_INTERVAL:
            gc.collect()
            chunks_since_gc = 0
        
        # Save checkpoint periodically
        chunks_since_checkpoint += 1
        if chunks_since_checkpoint >= CHECKPOINT_INTERVAL:
            checkpoint.shard_index = shard_index
            checkpoint.all_tokens_buffer = all_tokens_buffer
            checkpoint.total_docs_processed = total_docs_processed
            checkpoint.save(checkpoint_path)
            chunks_since_checkpoint = 0
        
        # Track total chunks for pool recreation
        total_chunks_processed += 1
        if total_chunks_processed % (MAX_TASKS_PER_WORKER * NUM_PROCESSES) == 0:
            print("\nRecreating worker pool to free memory...")
            pool = recreate_pool(pool, NUM_PROCESSES, initialize_worker, pool_init_args)
            # Recreate the results iterator with new pool
            results_iterator = pool.imap_unordered(
                process_json_lines_batch,
                line_chunk_generator,
                chunksize=POOL_DISPATCH_CHUNKSIZE
            )
            progress_bar.close()
            progress_bar = tqdm(results_iterator, total=num_chunks_estimate, desc=progress_desc)
    
    progress_bar.close()
    
    # Mark file as completed if we didn't stop mid-file
    if not shutdown_requested:
        checkpoint.completed_files.add(raw_file_path)
        checkpoint.current_file = None
        checkpoint.current_file_line = 0
        print(f"Completed processing {raw_file_path}")
        
        # Force GC after each file
        gc.collect()
finally: # Clean up pool pool.close() pool.join()

Save final checkpoint
checkpoint.shard_index = shard_index checkpoint.all_tokens_buffer = all_tokens_buffer checkpoint.total_docs_processed = total_docs_processed checkpoint.save(checkpoint_path)

Write remaining tokens
if len(all_tokens_buffer) > 0 and not shutdown_requested: # Convert deque to list for numpy remaining_tokens = list(all_tokens_buffer) tokens_np = np.array(remaining_tokens, dtype=np.uint16) split = "val" if shard_index == 0 else "train" output_filename = os.path.join(OUTPUT_BIN_DIR, f"mydata_{split}_{shard_index:06d}.bin") if write_datafile(output_filename, tokens_np): print(f"Wrote final shard {shard_index} with {len(remaining_tokens)} tokens") shard_index += 1

Final statistics
end_time = time.time() total_time = end_time - checkpoint.start_time print(f"\n{'Checkpoint saved - can resume later' if shutdown_requested else 'Processing completed!'}") print(f"Total documents processed: {total_docs_processed:,}") print(f"Total shards written: {shard_index}") print(f"Total time: {total_time/3600:.2f} hours") if total_docs_processed > 0: print(f"Processing speed: {total_docs_processed/total_time:.2f} docs/second")

if not shutdown_requested: # Clean up checkpoint file on successful completion checkpoint_files = [checkpoint_path, checkpoint_path + '.gz'] for cp_file in checkpoint_files: if os.path.exists(cp_file): os.remove(cp_file) print(f"Checkpoint file removed: {cp_file}") else: print(f"Resume by running the script again - checkpoint saved")

Final tokenizer info
final_tokenizer = Tokenizer.from_file(TOKENIZER_PATH) final_eot_id = final_tokenizer.token_to_id(EOT_TOKEN_STR_MAIN) print(f"\nIMPORTANT: Update train_gpt.py with:") print(f" - vocab_size: {final_tokenizer.get_vocab_size()}") print(f" - EOT token ID: {final_eot_id}") print(f" - Output directory: {OUTPUT_BIN_DIR}") the tokenization without these lines

import os os.environ["TOKENIZERS_PARALLELISM"] = "false"

Creates memory leaks, but with these lines it is much slowers how can we overcome the issue with memory and make a fast script