🤗 Dataset Train/Test Split

If you've ever hand-built a "train/test split" by taking every Nth row, this node is here to tell you that was a bad idea - and to hand you the correct tool. 🤗 Dataset Train/Test Split randomly divides a loaded dataset into two parts the way the datasets library intends, and because it's a node with a seed, the split is reproducible instead of a dice roll every time the graph runs.

How it works

It maps onto datasets.Dataset.train_test_split, so you get the real deal: a random partition (by default it shuffles first), with a controllable held-out fraction. The inputs are the whole story:

• test_size - fraction of rows held out for test, default 0.2 (20%). The step is 0.05, so it's a slider for intuition, but you can type a precise value.
• seed - the random seed. Set it and the same input always yields the same two piles. This is the reproducibility knob; if your downstream training run varies run to run, this is where the variance came from.
• shuffle - shuffle before splitting (default on). Turn it off only when you deliberately want the first 80% / last 20% of the current row order.

The split is stratified-ish by nothing - plain random, no class balancing. If your classes are lopsided and that matters, filter or order your rows first.

The two outputs

Unlike most nodes here, this one gives you two HUGGINGFACE_DATASET outputs, train and test, and that's the point: you wire them into separate downstream branches - one path for training, one for evaluation - and the graph keeps them apart cleanly. Each output is a normal dataset, so you can chain more 🤗 nodes onto either side.

Installing it

Same pack as every other 🤗 node here - 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, handled by Manager.

Read this before wiring it up

Two constraints bite people. First, this node is loaded-only: it needs a fully materialized datasets.Dataset, so if the loader feeding it has streaming on, you'll get a clear error telling you to disable streaming. Second, remember that a split is only as clean as its source - if you already filtered or deduplicated rows, the split respects that order, so do your shaping first and split last. And if your dataset already ships train/test splits on the Hub (most good ones do), you don't need this node at all - just load each split from the loader's dropdown. Reach for this when you loaded one flat table and need to carve it up yourself.