wikitext_dataset¶
- torchtune.datasets.wikitext_dataset(tokenizer: ModelTokenizer, source: str = 'EleutherAI/wikitext_document_level', subset: str = 'wikitext-103-v1', max_seq_len: Optional[int] = None, packed: bool = False, split: str = 'train', **load_dataset_kwargs: Dict[str, Any]) Union[TextCompletionDataset, PackedDataset] [source]¶
Support for family of datasets similar to wikitext, an unstructured text corpus consisting of fulls articles from Wikipedia.
- Parameters:
tokenizer (ModelTokenizer) – Tokenizer used by the model that implements the
tokenize_messages
method.source (str) – path to dataset repository on Hugging Face. For local datasets, define source as the data file type (e.g. “json”, “csv”, “text”) and pass in the filepath in
data_files
. See Hugging Face’sload_dataset
(https://huggingface.co/docs/datasets/en/package_reference/loading_methods#datasets.load_dataset.path) for more details.subset (str) – name of subset of data to use, see the wikitext page for available subsets. Default is
"wikitext-103-v1"
.max_seq_len (Optional[int]) – Maximum number of tokens in the returned input and label token id lists. Default is None, disabling truncation. We recommend setting this to the highest you can fit in memory and is supported by the model. For example, llama2-7B supports up to 4096 for sequence length.
packed (bool) – Whether or not to pack the dataset to
max_seq_len
prior to training. Default is False.split (str) –
split
argument fordatasets.load_dataset
. You can use this argument to load a subset of a given split, e.g.split="train[:10%]"
. Default is “train”.**load_dataset_kwargs (Dict[str, Any]) – additional keyword arguments to pass to
load_dataset
.
- Returns:
- the configured
TextCompletionDataset
or
PackedDataset
ifpacked=True
- the configured
- Return type:
Union[TextCompletionDataset, PackedDataset]