the_cauldron_dataset¶
- torchtune.datasets.multimodal.the_cauldron_dataset(model_transform: Transform, *, subset: str, source: str = 'HuggingFaceM4/the_cauldron', column_map: Optional[Dict[str, str]] = None, new_system_prompt: Optional[str] = None, packed: bool = False, split: str = 'train', **load_dataset_kwargs: Dict[str, Any]) SFTDataset [source]¶
Support for family of image + text datasets similar to The Cauldron from Hugging Face Datasets.
The Cauldron consists of numerous datasets. You must specify one of the datasets using the
subset
argument.The model transform is expected to be a callable that applies pre-processing steps specific to a model. For multimodal datasets, this is expected to be at minimum a tokenizer and an image transform. The tokenizer will convert text sequences into token IDs after the dataset is converted to a list of
Message
. The image transform will load the image and process it in accordance to the model’s requirements.Here is a minimal example for illustrative purposes:
from torchtune.models.llama3 import llama3_tokenizer from torchtune.models.clip import CLIPImageTransform from torchtune.modules.transforms import Transform class MyModelTransform(Transform): def __init__( self, tokenizer_path: str, max_seq_len: Optional[int] = None, ): self.tokenizer = llama3_tokenizer(tokenizer_path) self.image_transform = CLIPImageTransform() def __call__(self, sample: Mapping[str, Any]) -> Mapping[str, Any]: tokens, mask = self.tokenizer.tokenize_messages(sample["messages"]) images = self.image_transform(sample["images"]) return { "tokens": tokens, "mask": mask, "images": images, }
See
SFTDataset
for more details about model transforms and message transforms.- Parameters:
model_transform (Transform) – model-specific transform class that takes in a sample dict and applies custom transforms on the keys. It should consist of at minimum two components: text tokenization (called on the “messages” field) and image transform (called on the “images” field). The keys returned by the model transform should be aligned with the expected inputs into the model.
subset (str) – name of the subset of the dataset to load. See the dataset card for options.
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
for more details. Default isHuggingFaceM4/the_cauldron
.column_map (Optional[Dict[str, str]]) – a mapping to change the expected “images” and “texts” column names to the actual column names in the dataset. Default is None, keeping the default column names.
new_system_prompt (Optional[str]) – if specified, prepend a system message. 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
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
. See Hugging Face’s API ref for more details.
- Returns:
dataset configured with source data and transform
- Return type:
- Raises:
ValueError – If
packed
is True, they are not supported for multimodal datasets yet.
Example
>>> cauldron_ds = the_cauldron_dataset(model_transform=model_transform, subset="ai2d") >>> for batch in Dataloader(cauldron_ds, batch_size=8): >>> print(f"Batch size: {len(batch)}") >>> Batch size: 8