Add Dataset

Prerequisites

• Prefer repo-managed dependencies over ad hoc pip install.
• For repeated work in a checked-out Marin repo, install the synced environment with uv sync --all-packages.
• For one-off schema inspection, run the tool with ephemeral deps via uv run --with datasets --with pyyaml lib/marin/tools/get_hf_dataset_schema.py ....
• Ensure access to the dataset: Hugging Face Hub ID, local path, or other supported formats.
• Use uv run lib/marin/tools/get_hf_dataset_schema.py from a synced Marin checkout, or the uv run --with ... form above when the environment is not already provisioned.

Guidelines for Humans

Command Line Usage

uv run lib/marin/tools/get_hf_dataset_schema.py <dataset_name> [options]

Python Import

from marin.tools.get_hf_dataset_schema import get_schema
schema = get_schema(dataset_name="wikitext", config_name="wikitext-103-v1")

Common Workflows

1. Basic inspection: uv run lib/marin/tools/get_hf_dataset_schema.py roneneldan/TinyStories
2. Multi-config dataset: First try without config, then use the suggested config from the error
3. Remote code datasets: Add --trust_remote_code flag when needed

Rules for Agents

1. Config Handling

• ALWAYS check if a dataset requires a config first
• If configs are required, the tool returns:

  {
    "error": "Config name is required.",
    "available_configs": ["config1", "config2", ...]
  }
  

• Select an appropriate config based on the available options

2. Text Field Selection

• Prioritize fields named exactly 'text'
• Fall back to fields containing 'text' in their name
• Consider string-type fields if no obvious text field exists
• Examine sample_row to verify field contents

3. Error Handling

• Handle common error cases:
  • Missing config
  • Dataset not found
  • Remote code execution required
• Retry with appropriate parameters (e.g., --trust_remote_code)

4. Performance

• Tool uses streaming to avoid downloading full datasets
• Expect quick responses for schema info
• Sample row might be empty for some datasets

Output Format

The tool returns a JSON object with:

{
  "splits": ["train", "validation", ...],
  "text_field_candidates": ["text", "content", ...],
  "features": {
    "text": "string",
    "label": "int64",
    ...
  },
  "sample_row": {
    "text": "Example content...",
    ...
  }
}

Examples

Basic Usage

$ uv run lib/marin/tools/get_hf_dataset_schema.py roneneldan/TinyStories
{
  "splits": ["train", "validation"],
  "text_field_candidates": ["text"],
  "features": {"text": "string"},
  "sample_row": {"text": "Once upon a time..."}
}

Dataset with Config

$ uv run lib/marin/tools/get_hf_dataset_schema.py wikitext
{
  "error": "Config name is required.",
  "available_configs": ["wikitext-103-raw-v1", "wikitext-103-v1", ...]
}

$ uv run lib/marin/tools/get_hf_dataset_schema.py wikitext --config_name wikitext-103-v1
{
  "splits": ["train", "validation", "test"],
  "text_field_candidates": ["text"],
  "features": {"text": "string"},
  "sample_row": {"text": "Article content..."}
}

Dataset with Remote Code

$ uv run lib/marin/tools/get_hf_dataset_schema.py c4 --config_name en --trust_remote_code
{
  "splits": ["train", "validation"],
  "text_field_candidates": ["text"],
  "features": {"text": "string", "url": "string", "timestamp": "string"},
  "sample_row": {"text": "Web content..."}
}

Next Steps

Once the schema is inspected and the dataset is registered (for example in experiments/pretraining_datasets/__init__.py or in a dedicated file), the goal is to cargo-cult existing dataset configs for tokenization:

• Apply transformations (e.g., field mapping).
• Estimate token counts and file sizes.
• Find similar dataset configurations in Marin's existing experiments.
• Copy and adapt tokenization configs from similar datasets.
• Run ablations or trials.

See Also

• lib/marin/tools/get_hf_dataset_schema.py
• Hugging Face datasets documentation