Method __init__

codegeex/megatron/data/prompt_dataset.py:180–227

Args: name: name of the dataset. data_prefix: prefix of the data. documents: list of document indices. input_ids_indexed_dataset: indexed dataset for prompts. attention_mask_index_dataset: indexed dataset for text. labe

(
        self,
        name,
        data_prefix,
        documents,
        input_ids_indexed_dataset,
        attention_mask_index_dataset,
        labels_indexed_dataset,
        num_samples,
        seq_length,
        seed,
    )

class PromptDataset(torch.utils.data.Dataset):
    def __init__(
        self,
        name,
        data_prefix,
        documents,
        input_ids_indexed_dataset,
        attention_mask_index_dataset,
        labels_indexed_dataset,
        num_samples,
        seq_length,
        seed,
    ):
        """
        Args:
            name: name of the dataset.
            data_prefix: prefix of the data.
            documents: list of document indices.
            input_ids_indexed_dataset: indexed dataset for prompts.
            attention_mask_index_dataset: indexed dataset for text.
            labels_indexed_dataset: indexed dataset for labels.
            num_samples: number of samples to draw from the indexed dataset.
            seq_length: sequence length.
            seed: seed for random number generator.
        """

        self.name = name
        self.input_ids_indexed_dataset = input_ids_indexed_dataset
        self.attention_mask_index_dataset = attention_mask_index_dataset
        self.labels_indexed_dataset = labels_indexed_dataset
        self.seq_length = seq_length
        self.eod_token = get_tokenizer().eod

        # Checks
        assert np.min(documents) >= 0
        assert np.max(documents) < input_ids_indexed_dataset.sizes.shape[0]
        assert input_ids_indexed_dataset.sizes.shape[0] == attention_mask_index_dataset.sizes.shape[0]
        assert attention_mask_index_dataset.sizes.shape[0] == labels_indexed_dataset.sizes.shape[0]

        # Build index mappings.
        self.doc_idx = _build_index_mappings(
            self.name,
            data_prefix,
            documents,
            self.input_ids_indexed_dataset.sizes,
            num_samples,
            seq_length,
            seed,
        )

    def __len__(self):
        # -1 is due to data structure used to retieve the index: