🤗 Dataset Take

You almost never want all fifty million rows. You want a few thousand to test a pipeline, a hundred to eyeball, ten to debug a node. 🤗 Dataset Take is the "keep the first N and drop the rest" node, and it's the single most useful thing to drop in front of a big streaming load - it's how you make a lazy dataset actually lazy instead of a download-the-whole-world event.

How it works

It's datasets.Dataset.take (and IterableDataset.take) as a node. Two inputs:

• dataset - what to trim.
• n - how many leading rows to keep (default 100).

One output: a dataset limited to its first n rows. That's the whole node, and it's worth more than it looks because of when it matters.

The streaming story is the real story

With the loader's streaming on, Take is what keeps your graph fast. Feed a lazy IterableDataset into Take(50), then into 🤗 Dataset To Data List / To LIST, and only those fifty rows ever get fetched from the Hub - the rest of the corpus is never downloaded. The README's own example workflow (imdb_streaming_skip_take) is exactly this pattern. Without a Take (or the loader's limit widget), even a "small" materialization of a streaming dataset can pull far more than you meant.

On a fully-loaded dataset Take is simpler: it's just a cap, a cheap way to bound how many rows your downstream loop processes. Pair it with Skip and you get a middle window (Skip(500) → Take(100)); pair it with Shuffle and you get a pseudo-random sample; pair it with Sort and you get "top N by column" - sort first, then take.

Gotchas

Take keeps the first N in the current order, whatever that order is - unshuffled Hub data is usually in whatever order the author pushed it, so if you want a representative slice, shuffle (with a seed) before you take. And remember Take is positional, not selective: if you want specific rows, that's Select Rows' job (loaded-only); Take is for the front of the stream, which is why it's the streaming-safe one.

Installing it

Part of StableLlama/ComfyUI-huggingface_dataset. ComfyUI-Manager → search "Hugging Face dataset", or:

cd ComfyUI/custom_nodes
git clone https://github.com/StableLlama/ComfyUI-huggingface_dataset
pip install -r requirements.txt

Restart ComfyUI. Only runtime dependency is datasets, installed for you by Manager.