Getting your data

Models need data to train on. Data can be either curated using proprietary and/or readily available sources, or synthetically generated by the environment a model trains in, potentially with ground truth data to verify against.

Here, we’ll focus on the former, and in the next guide, we’ll show how you can do the latter.

Hugging Face

The HuggingFaceDataset class handles all the subtle ways in which datasets differ.

For example, statworx/haiku contains a keywords column with only a single word as input, and a text column that contains a ground-truth label. Since the input isn’t in the OpenAI chat completions API format, we need to set apply_chat_template to True. By default, the Gym will format each prompt in the dataset as a single user message.

from modal_training_gym import HuggingFaceDataset


class HaikuDataset(HuggingFaceDataset):
    hf_repo = "statworx/haiku"
    input_column = "keywords"
    output_column = "text"
    apply_chat_template = True

You can customize this by adding a system prompt, or by providing a prompt template to add additional text to the user message:

class HaikuDataset(HuggingFaceDataset):
    # ...
    apply_chat_template = True
    system_prompt = "You are an expert poet."
    prompt_template = "Write a haiku about {input}."

Note that only input is allowed.

Other datasets like zhuzilin/dapo-math-17k are already formatted using a chat template:

class MathDataset(HuggingFaceDataset):
    hf_repo = "zhuzilin/dapo-math-17k"
    input_key = "prompt"
    label_key = "label"
    apply_chat_template = False

As any seasoned ML veteran will tell you, we need separate datasets for training and validation/evaluation to properly train a model. This is made easy with HF’s slicing syntax:

train_dataset = MyHFDataset(..., hf_split="train[:1000]")
eval_dataset = MyHFDataset(..., hf_split="train[1000:]")

If a dataset is gated, you can create a Modal Secret named huggingface-secret that the Gym will auto-detect:

modal secret create huggingface-secret HF_TOKEN=hf_...

Harbor

A Harbor dataset provides a series of tasks, each containing instructions, files for the coding environment, and tests.

from modal_training_gym import HarborDataset


class HelloWorld(HarborDataset):
    dataset_name="harbor/hello-world"

If you have the files installed locally, you can instead use the filepath:

class HelloWorld(HarborDataset):
    path="/path/to/task_root"

Some tasks provide additional metadata:

class HelloWorld(HarborDataset):
    # ...
    label_metadata_path="task.toml"
    test_data_dir="tests"

During training, each task will be converted into a single user prompt. Like HuggingFaceDataset, you can optionally provide a system prompt or override the template for the user message. However, if the metadata contains variables, you can also use them in prompt_template!

class HelloWorld(HarborDataset):
    # ...
    system_prompt="You are an expert Python programmer."
    prompt_template="Perform task {task_name} stored at {task_path}: {instruction}"

Creating a custom dataset

To use your own data, likely stored in an external source or a Modal Volume, you simply subclass DatasetConfig:

import os

from datasets import Dataset
from modal_training_gym import DatasetConfig

prompts = [(prompt, label), ...]  # external source


class MyCustomDataset(DatasetConfig):
    input_key = "messages"
    label_key = "label"
    apply_chat_template = True

    def prepare(self, path: str, eval_paths: dict[str, str] | None = None) -> None:
        os.makedirs(os.path.dirname(path), exist_ok=True)
        rows = [
            {
                self.input_key: [{"role": "user", "content": prompt}],
                self.label_key: label,
            }
            for prompt, label in prompts
        ]
        Dataset.from_list(rows).to_parquet(path)
        if eval_paths:
            for eval_path in eval_paths.values():
                os.makedirs(os.path.dirname(eval_path), exist_ok=True)
                Dataset.from_list(rows).to_parquet(eval_path)


dataset = MyCustomDataset()