Is this page helpful?

core.datasets.gpt_fim_dataset#

Module Contents#

Classes#

GPTFIMDatasetConfig

Configuration object for Megatron Core GPT FIM datasets

GPTFIMDataset

The GPT dataset with fill-in-the-middle preprocessing

Data#

logger

_FIM_TOKEN_KEYS

API#

core.datasets.gpt_fim_dataset.logger#

‘getLogger(…)’

core.datasets.gpt_fim_dataset._FIM_TOKEN_KEYS#

(‘prefix’, ‘middle’, ‘suffix’, ‘pad’, ‘eod’)

class core.datasets.gpt_fim_dataset.GPTFIMDatasetConfig#

Bases: megatron.core.datasets.gpt_dataset.GPTDatasetConfig

Configuration object for Megatron Core GPT FIM datasets

fim_rate: float#

None

Probability to convert a training sample into a FIM format

fim_spm_rate: float#

None

Probability that the a FIM sample uses the SPM format over the PSM format

fim_extra_tokens: Dict#

None

FIM extra tokens. Should consist of prefix, middle, suffix, PAD, and EOD tokens.

fim_split_sample: Optional[str]#

None

String around which to split the sample for FIM

fim_fragment_rate: Optional[float]#

None

Rate of FIM on each fragment when split_sample is not None

fim_no_prefix: Optional[str]#

None

Do not apply FIM to fragments that start with this prefix

validate_without_tokenizer() → None#

The GPT checks that need no tokenizer, then the FIM ones, which need none either: both rates and the five extra tokens must be given (a rate of 0.0 is a value; validate_args checks the same for --fim-data). validate() runs this before the tokenizer checks.

class core.datasets.gpt_fim_dataset.GPTFIMDataset(
indexed_dataset: megatron.core.datasets.indexed_dataset.IndexedDataset,
dataset_path: str,
indexed_indices: numpy.ndarray,
num_samples: int,
index_split: megatron.core.datasets.utils.Split,
config: core.datasets.gpt_fim_dataset.GPTFIMDatasetConfig,
)#

Bases: megatron.core.datasets.gpt_dataset.GPTDataset

The GPT dataset with fill-in-the-middle preprocessing

Parameters:

  • indexed_dataset (IndexedDataset) – The IndexedDataset around which to build the

  • MegatronDataset

  • indexed_indices (np.ndarray) – The set of the documents indices to expose

  • num_samples (int) – The number of samples to draw from the indexed dataset

  • index_split (Split) – The indexed_indices Split

  • config (GPTFIMDatasetConfig) – The GPT-specific container for all config sourced parameters

Initialization

_query_document_sample_shuffle_indices(
idx: int,
) → Tuple[numpy.ndarray, numpy.ndarray, list]#

Get the text (token ids) and document ids for a given index

Parameters:

idx (int) – The index into the dataset

Returns:

The text ids, document ids, and per-document token counts.

Return type:

Tuple[np.ndarray, np.ndarray, list]

_fim_permute_sequence(sequence, rate)#
_fim_split_and_permute_sequence(sequence)#

If self.fim_split_sample is not None, split the sequence. Then apply FIM on the fragments, or the whole sequence if self.fim_split_sample is None.

_permute(
sample,
fim_rate,
fim_spm_rate,
tokenizer,
truncate_or_pad=True,
suffix_tok_id=None,
prefix_tok_id=None,
middle_tok_id=None,
pad_tok_id=None,
no_fim_prefix=None,
)#

Take in a sample (np array w/ size (0,chunklength)) and perform a FIM transformation on it. Maintain the same sample length (if transform creates a few extra tokens, drop them).