Function process_files

scripts/data_preprocess.py:10–79

Process a specified number of lines from each .jsonl.zst file in the input directory and save encoded tokens to an HDF5 file. Args: input_dir (str): Directory containing input .jsonl.zst files. output_file (str): Path to the output HDF5 file. tokenizer_name (str

(input_dir: str, output_file: str, tokenizer_name: str, max_data: Optional[int] = None)

from typing import Optional

def process_files(input_dir: str, output_file: str, tokenizer_name: str, max_data: Optional[int] = None) -> None:
    """
    Process a specified number of lines from each .jsonl.zst file in the input directory
    and save encoded tokens to an HDF5 file.

    Args:
        input_dir (str): Directory containing input .jsonl.zst files.
        output_file (str): Path to the output HDF5 file.
        tokenizer_name (str): Name of the tiktoken tokenizer to use (e.g., 'r50k_base').
        max_data (int, optional): Maximum number of lines to process from each file.
                                  If None, process all lines.
    """
    # Print processing strategy based on max_data
    if max_data is not None:
        print(f"You have chosen max_data = {max_data}. Processing only the top {max_data} JSON objects from each file.")
    else:
        print("Processing all available JSON objects from each file.")

    # Load the tokenizer using the provided tokenizer name
    enc = tiktoken.get_encoding(tokenizer_name)

    # Create an HDF5 file for output
    with h5py.File(output_file, 'w') as out_f:
        # Initialize the dataset for storing tokenized data
        dataset = out_f.create_dataset('tokens', (0,), maxshape=(None,), dtype='i')
        start_index = 0  # Track the starting index for the next batch of tokens

        # Process each .jsonl.zst file in the input directory
        for filename in sorted(os.listdir(input_dir)):
            if filename.endswith(".jsonl.zst"):  # Only process .jsonl.zst files
                in_file = os.path.join(input_dir, filename)
                print(f"Processing: {in_file}")

                processed_lines = 0  # Counter for processed lines in the current file

                # Open the compressed .jsonl.zst file for reading
                with zstd.open(in_file, 'rt', encoding='utf-8') as in_f:
                    # Iterate over each line in the file
                    for line in tqdm(in_f, desc=f"Processing {filename}", total=max_data if max_data is not None else None):
                        try:
                            # Parse the line as JSON
                            data = json.loads(line)
                            text = data.get('text')  # Extract the 'text' field from the JSON object

                            if text:
                                # Tokenize the text and append an end-of-text token
                                encoded = enc.encode(text + "<|endoftext|>", allowed_special={'<|endoftext|>'})
                                encoded_len = len(encoded)

                                # Resize the dataset to accommodate new tokens
                                end_index = start_index + encoded_len
                                dataset.resize(dataset.shape[0] + encoded_len, axis=0)

                                # Store the encoded tokens in the dataset
                                dataset[start_index:end_index] = encoded
                                start_index = end_index  # Update the start index
                            else:
                                # Warn if 'text' key is missing in the JSON object