samsum_dataset

torchtune.datasets.samsum_dataset(tokenizer: ModelTokenizer, *, source: str = 'Samsung/samsum', column_map: Optional[Dict[str, str]] = None, train_on_input: bool = False, new_system_prompt: Optional[str] = None, packed: bool = False, split: str = 'train', **load_dataset_kwargs: Dict[str, Any]) → Union[SFTDataset, PackedDataset]

Support for summarization datasets and their variants from Hugging Face Datasets. An example is the SAMsum dataset.

It is recommended to configure the tokenizer with the SummarizeTemplate in conjunction with this dataset.

Masking of the prompt during training is controlled by the train_on_input flag, which is set to False by default - If train_on_input is True, the prompt is used during training and contributes to the loss. - If train_on_input is False, the prompt is masked out (tokens replaced with -100)

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”), pass in the filepath in data_files, and set split="train". See Hugging Face’s load_dataset for more details. Default is Samsung/samsum.

• column_map (Optional[Dict[str, str]]) – a mapping from the expected columns in the message transform InputOutputToMessages to the new column names in the dataset. Keys should be “input” and “output” and values should be the actual column names. If None, use the default column names {"input": "dialogue", "output": "summary"} in Samsung/samsum.

• train_on_input (bool) – Whether the model is trained on the prompt or not. Default is False.

• new_system_prompt (Optional[str]) – if specified, prepend a system message to every sample. This can serve as instructions to guide the model response. Setting this will OVERRIDE any system messages already present in the dataset. Default is None.

• packed (bool) – Whether or not to pack the dataset to tokenizer’s max_seq_len prior to training. Default is False.

• split (str) – split argument for datasets.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:

dataset configured with source data and template

Return type:

Union[SFTDataset, PackedDataset]

Raises:

ValueError – If packed=True and tokenizer.max_seq_len is not set.

Example

>>> samsum_ds = samsum_dataset(model_transform=tokenizer)
>>> for batch in Dataloader(samsum_ds, batch_size=8):
>>>     print(f"Batch size: {len(batch)}")
>>> Batch size: 8