Awesome

Generative data preparation

This software package allows you to prepare datasets for training generative LLMs on SambaStudio and SambaNova's Reconfigurable Data Units (RDUs). Some features include efficient multiprocessing, shuffling data that outsizes RAM, and specifying tokens to attend to during training.

The pipeline.py script streamlines the data preparation process. It takes a single input file, shuffles and splits it into train/dev/test files, tokenizes, sequences, and converts them to HDF5 format using the utilities in data_prep.py. The output directory contains multiple split HDF5 files that are needed to run data parallel training. This output directory will be directly used as a training dataset in SambaStudio. While this package features simple flows that work out of the box, it also supports more customization allowing for many styles of packing varied length text into tokenized sequences.

If you are an advanced user looking to process data with pre-defined splits, integrate with the package validation tools, or contribute, check out the Advanced Usage section below!

Table of contents

• Requirements
• Installation
• Getting Started
• Input
• Formatting data for Chat/Instruction/Fine Tuned Models
• Output
• Flags
• Examples
  • Pre-training
  • Fine-tuning
  • Dialogue
  • Meta in context learning
• Understanding Command Outputs
• FAQs
• Advanced Usage

Requirements

• Python version 3.8.10+
• Support for Linux and Mac OS. Not tested on Windows

Installation

git clone https://github.com/sambanova/generative_data_prep.git
cd generative_data_prep
pip install .

Getting Started

The following simple example will help you get started with your first processed dataset:

Example

python3 -m generative_data_prep pipeline --input_path=<PATH TO DATASET FILE> --output_path=<PATH TO OUTPUT DIRECTORY> --pretrained_tokenizer=openai-community/gpt2 --max_seq_length=1024 --input_packing_config='greedy::drop' --shuffle=on_RAM

Here are a few important parameters to know about when running this example:

Input

The input_path argument must be a file or a directory containing one files, each file must be a .txt or .jsonl.

.jsonl Format

The JSON Lines format can be used for fine-tuning, or pre-training/continual pre-training. Each line in the .jsonl format should be a json object with a prompt, and completion element. For example:

{"prompt": "What did the fox do?", "completion": "The quick brown fox jumped over the lazy dog."}
{"prompt": "How much wood does a woodchuck chuck?", "completion": "A woodchuck chucks 1000 wood."}
{"prompt": "Who sells seashells by the sea shore?", "completion": "She sells seashells by the sea shore."}

We also support lists of prompt/completion pairs within a .jsonl file. This guarantees that the prompt/completion pairs in the list will be placed contiguously in the same sequence. If the input prompt/completion pairs are placed on separate lines rather then a list, then they will get shuffled and appear in different training sequences. Your input file may include lines in both list format and regular prompt/completion pair format. Here's an example structure:

[{"prompt": "What's your favorite type of music?", "completion": "I love hip-hop"}, {"prompt": "That's cool. Who's your favorite rapper?", "completion": "I really like Kendrick Lamar"}]
[{"prompt": "What is your favorite type of dessert?", "completion": "My favorite dessert is cheesecake."}, {"prompt": "What is your favorite flavor of cheesecake?", "completion": "My favorite flavor of cheesecake is raspberry."}]
[{"prompt": "What is your favorite sport?", "completion": "My favorite sport is football."}, {"prompt": "Who is your favorite football player?", "completion": "My favorite football player is Tom Brady."}]

If the JSON objects in your .jsonl contain keywords other than prompt and completion, refer to the prompt_keyword and completion_keyword flags below

.txt Format

This format should only be used for pre-training/continual pre-training, but not fine-tuning. Additionally, even though .txt format is supported, we recommend that you still use prompt/completion .jsonl format because it can handle newlines in the text. If you use .txt format, then newlines within individual text articles will seperate the text into different data points that may be shuffled and not placed into the same contiguous sequences.

The quick brown fox jumped over the lazy dog
I come from a land down under
SambaNova makes extremely good software and hardware that's fun to use

The above txt input would be equivalent to this jsonl input

{"prompt": "", "completion": "The quick brown fox jumped over the lazy dog"}
{"prompt": "", "completion": "I come from a land down under"}
{"prompt": "", "completion": "SambaNova makes extremely good software and hardware that's fun to use"}

Formatting data for Chat/Instruction/Fine Tuned Models

Many chat and instruct models require very specific formatting to input multi turn conversations for training and inference. The tokenizer.apply_chat_template function easily adapts your jsonl data to this format. To use this feature, prepare your data in jsonl format as specified above, and then include the --apply_chat_template flag to automatically prepare your data in this format.

If your data is in the classic chat template format like [{"role": "user", "content": "..."}...], and you would like to convert it into the prompt completion format to be compatible with this repo, please use the generative_data_prep/utils/convert_chat_template_to_prompt_completion.py script.

Output

The output_path should be a directory that will contain all the tokenized HDF5 split files, and a sub-directory called tokenizer. This output directory constitutes a processed dataset and can be used for training a model after uploading to SambaStudio. The tokenizer sub-directory will be transferred to any output checkpoints that are saved by Sambastudio for the tokenizer to be used for inference later on.

Holdout Evaluation and Test Data

To evaluate on a holdout set of data during training, pipeline.py can create splits of holdout evaluation and test data.

To do this, choose only one of the two options below. Please review the Flags section for in detail descriptions of these flags.

• To specify the number of training splits and evaluation splits directly, use the three flags --num_training_splits=..., --num_dev_splits=... and --num_test_splits=...

OR

• To specify the percentage of the data heldout for evaluation, you can specify --dev_ratio=0.1 and --test_ratio=..., where 0.1 means that approximately 10% of the data will be included in the evaluation splits. You can also specify the --num_training_splits=... flag to control the total number of training splits, but we recommend to let this default.

All this evaluation data will saved under the <OUTPUT_DIR>, if you want to run evaluation on the eval_splits during training you must enable do_eval on SambaStudio. All test data will be saved under <OUTPUT_DIR>/test. This data is left in jsonl text format because running evaluation or inference usually requires text inputs instead of tokenized inputs.

View Decoded HDF5 Files

If you want to view the contents of a processed dataset, you can decode an HDF5 file into a human readable text format. To do so, run the following command:

python3 generative_data_prep/utils/decode_hdf5.py --pretrained_tokenizer=<HF TOKENIZER KEY> --hdf5_file_path=<PATH TO HDF5 FILE> --output_decoded_file_path=<PATH TO OUTPUT TXT FILE>

Note: The same tokenizer used to prepare the data must be used for decoding!

Dataset Size Requirements

1. You need to ensure your dataset is large enough to run one batch of training.
2. Make sure that the number of sequences in the output dataset files satisfy this by checking max_batch_size_train in the <OUTPUT_DIR>/metadata.yaml file.
3. Ensure that the batch_size hyper-parameter is <= max_batch_size_train during training. To understand more, expand the details section below or see FAQs

How to Check and Set

When starting a training job, ensure that the batch_size hyper-parameter is no bigger than the max_batch_size_train shown in metadata.yaml.

For example:

$ cat <PROCESSED DATA DIRECTORY>/metadata.yaml

max_batch_size_dev: null
max_batch_size_train: 7
max_seq_length: 1024
number_of_dev_files: 0
number_of_test_files: 0
number_of_training_files: 32
token_type_ids: true
tokenizer_model_type: <class 'transformers.models.gpt2.configuration_gpt2.GPT2Config'>
vocab_size: 50257

Here you can see that max_batch_size_train is 7, so the batch size hyper-parameter cannot be greater than 7.

Explanation

With a sufficiently large dataset, you are generally fine with the defaults and can ignore. However, when the provided dataset is small (~1000 data points or less), you need to set the above values correctly or else you will likely run into a training error.

<br /> The dataset that you are providing will be split up across multiple hdf5 files based on the input parameters of the pipeline command.

• max_seq_length - The maximum sequence length the model you are using can take for a single data point. See more in flags section.
• input_packing_config - Determines how to pack the provided data into sequences that will be split across the hdf5 files for training. See more in the flags section.

Based on the size and strucutre of the dataset provided + these parameter settings, a different max_batch_size_train will be shown in metadata.yaml which dictates how large you can set the corresponding batch_size hyper-parameter setting when starting a model training job!

Note:: Not all models trained in studio will expose the batch_size parameter. For those that don't you should ensure your max_batch_size_train is larger than the default batch size (generally 16).

Additional Details

If you include the keep_split_jsonls flag, then the output_path will additionally contain a splits directory that saves the jsonl versions of the HDF5 files, meaning that splits/train_1_of_X.jsonl is the jsonl text version of train_1_of_X.hdf5.

The output HDF5 files each contain two datasets:

• input_ids: sequences of tokens ids
• token_type_ids: describe the type of each token. The default id assignments are:
  • id=0 for tokens in the prompt
  • id=1 for tokens in the completion
  • id=2 for <eos> tokens that serve as padding tokens (will not be trained to predict)
  • id=3 for <eos> tokens at the end of articles, that define the attention boundary when training with article attention

Flags

This section outlines all the flags you can set to customize the data prep pipeline for your use case!

Examples

Fine-tuning

Fine-tuning (also known as "generative tuning") is a technique used to adapt a pre-trained language model to perform better at a specific task. This approach typically involves training the model on input data that is structured as a "prompt" followed by a "completion". The prompt represents the input for a specific task, while the completion is the output that the model should generate. During training, the model learns to generate the relevant completion tokens based on the context provided by the prompt tokens.

The benefit of using this training format is that the model can learn to generate high-quality outputs for a specific task without requiring a large amount of task-specific training data. By leveraging the pre-trained language model's knowledge gained from being trained on a large corpus of text data, the fine-tuned model can quickly adapt to the new task and generate high-quality outputs with minimal training data.

When training on this kind of data using SambaStudio, set prompt_loss_weight=0.0. This ensures that the model does not learn to generate the prompt tokens, and only learns to generated completion tokens.

Example data

For fine-tuning, your data should be in .jsonl format with prompts and completions designed for the task you're adapting to.

Jsonlines with a prompt and completion

Example command

python3 -m generative_data_prep pipeline --input_path=./tests/examples/generative_tuning/example_generative_tuning_data.jsonl --output_path=./tests/examples/generative_tuning/pipelined_generative_tuning --pretrained_tokenizer=openai-community/gpt2 --max_seq_length=1024 --shuffle=on_RAM --input_packing_config=single::drop

View decoded output

Pre-training

Pre-training on unstructured data enables large languages models to learn general language patterns and structures that are useful for a wide range of downstream tasks. In order to prepare pre-training data, you need a large amount of unstructured text data. To prepare pre-training data use the flag --input_packing_config=full.

Example data

For pre-training you can have your data in two formats.

text separated by newlines.

jsonlines with empty prompts and text in the completions.

We recommend to use jsonlines with empty prompts and all the text in the completion, this is so that newlines in the text do not separate semantically related articles.

Example command

python3 -m generative_data_prep pipeline --input_path=./tests/examples/pretraining/example_pretraining_data.jsonl --output_path=./tests/examples/pretraining/pipelined_pretraining --pretrained_tokenizer=openai-community/gpt2 --max_seq_length=1024 --shuffle=on_RAM --input_packing_config=full

View decoded output

Dialogue

Dialogue data often involves multiple turns in a conversation between a user and an agent. In order to train on this data, the entire conversation needs to be in the same sequence of tokens and the model should only learn to generate the agents responses based on the users inputs. To prepare data like this create a list of prompt completion pairs, and if you train with packing_boundary=jsonl and input_packing_config=greedy::truncate_right/ or input_packing_config=single::truncate_right then these conversations are guaranteed to be in the provided order in the same sequence. Additionally if you include the prompt_loss_weight=0.0 option while training on SambaStudio, only the completions will be learned. Also for training dialogue in chat-ml style, users can set prompt_prefix and prompt_postfix.

Example data

Lists of prompt completion pairs that represent turns in a conversation

Example command

python3 -m generative_data_prep pipeline --input_path=./tests/examples/dialogue/example_dialogue_data.jsonl --output_path=./tests/examples/dialogue/pipelined_dialogue --pretrained_tokenizer=openai-community/gpt2 --max_seq_length=1024 --shuffle=on_RAM --input_packing_config=single::truncate_right

View decoded output

Meta in context learning

Meta In Context Learning improves the few shot performance of a model by including training data formatted in a few shot style. This infrastructure allows you to prepare data in a variant of meta in context learning SambaNova uses called "All Shot" learning. In order to prepare data in this format prepare lists of prompt completion pairs, where every list contains prompt completion pairs that are completing the same instruction/task. Then prepare the data with the input_packing_config=greedy::drop, packing_boundary=prompt_completion_pair and attention_boundary=jsonl. This ensures that every sequence contains prompt completion pairs following the same "instruction", and that when learning a completion the model is attending to all the other prompt completion pairs before it.

Example data

Lists of prompt completion pairs that are all from the same task

Example command

python3 -m generative_data_prep pipeline --input_path=./tests/examples/metaICL/example_metaICL_data.jsonl --output_path=./tests/examples/metaICL/pipelined_metaICL --pretrained_tokenizer=openai-community/gpt2 --max_seq_length=1024 --shuffle=on_RAM --input_packing_config=greedy::drop --packing_boundary=prompt_completion_pair --attention_boundary=jsonl

View decoded output

Understanding Command Outputs

Terminal Output

The metrics associated with this dataset will be printed in the terminal as well as being logged at <OUTPUT DIR PATH>/logs.log. These metrics give some insight into how the data was packed into sequences, and information about the training dataset.

Metadata Output File

To help improve speed and cross-checking we provide a metadata file along with the dataset. This file is located right under the output_dir as metadata.yaml. This file is used along with a custom pydantic model which you can import from this library which will verify the dataset parameters and the training parameters. This can be used as a way to catch bugs before training begins.

max_seq_length: int
token_type_ids: bool
vocab_size: int
tokenizer_model_type: str
number_of_training_files: int
number_of_dev_files: int
number_of_test_files: int
max_batch_size_train: int
max_batch_size_dev: Optional[int]

NOTE:

• tokenizer_model_type is the string conversion of type(modelConfig). Can use this field to compare the model used during training, which can be extracted by using AutoConfig in Huggingface transformers. Then wrapping it with str(type()).
• max_batch_size_dev will be None unless dev files are created during generative data pipeline.
• token_type_ids will always be True for now since they are always generated.

FAQs

"Cannot access gated repo" error

If you pass in a --pretrained_tokenizer for a model tokenizer that is gated on Huggingface, you need to get access to the model on HuggingFace by going to the model card and requesting access, then follow this documentation to generate a HuggingFace API key and finally log in on the HuggingFace CLI.

If you have the model checkpoint downloaded locally you can also pass in the path to the model checkpoint as the --pretrained_tokenizer!

Number of samples in each file must be greater than or equal to batch size

This error will occur if you try to run training with a batch size that is greater than the maximum batch size of the prepared dataset. The maximum batch size is printed in the terminal as "Batch size <=..." and also logged in the logs.log file in the output directory.

To fix this, you can do one of the following:

1. Increase the amount of input data you use.
2. Change to a "single" input packing configuration like single::truncate_right, which will not pack the sequences with multple data points, and therefore create more training sequences. However, this may cause training to be inefficient because a lot of the available sequence length is wasted with padding tokens.
3. Decrease the num_training_splits so that each split has more data. Keep in mind, however, that you must have more training splits than the number of parallel RDUs you use to train.

Advanced Usage

The following are some advanced usage patterns that may be applicable to you. Follow the links for more information:

• If you have data that has been custom pre-split, and you would like to tokenize these files individually, check out the Single File Tokenization Guide
• If you want to build in custom dataset validation with our pydantic model, look at our section on pydantic dataset validation.
• If you want to build in dataset verification checks, look at our section on checking for dataset corruption.
• If you want to contribute to this project, check out the contribution section.