Nodes/ComfyUI-Pt-Wrapper/Ptv Hf Dataset With Token Encode
ComfyUI Node

Ptv Hf Dataset With Token Encode

Text classification data, straight from Hugging Face

By HowToSD·Created about a year ago·Updated about a year ago· 7
Ptv Hf Dataset With Token Encode
  • encode
  • PTVDATASET
name
splittrain
sample_field_nametext
label_field_namelabel
remove_html_tagsfalse
encode_return_dictfalse

This is the data-loading half of the pack's text-classification story. Want to train (or fine-tune) a model on a Hugging Face dataset like IMDB sentiment without writing a line of Python? You string together this node, a tokenizer node, a model, and a trainer - and this one fetches the dataset and converts every review into token IDs your model can actually chew on.

How it works

It loads a Hugging Face dataset by name with the datasets library, then runs an encode callable over each text sample on the fly. That encode function is a PTCALLABLE - you produce it from the pack's Hf Tokenizer Encode node (Hugging Face tokenizer) or Sp Encode node (SentencePiece). For each row it returns the token IDs and an attention mask as tensors, plus the label as a long tensor. The squeeze(0) in the code tells you the tokenizers are expected to return batch-shaped tensors, which they do when padding/truncation is on.

The inputs that matter

  • name - the HF dataset id, e.g. stanfordnlp/imdb.
  • split - default train. Use test (or whatever the dataset calls it) for evaluation.
  • sample_field_name / label_field_name - defaults text and label. If your dataset calls them review and sentiment, change these.
  • encode - the PTCALLABLE from a tokenizer node. Not optional; this is the node's whole point.
  • remove_html_tags - strips HTML before tokenizing. Nice for scraped datasets.
  • encode_return_dict - here's the one that bites people. The HF tokenizer encode returns a dict ({"input_ids": ..., "attention_mask": ...}), so set this to True. The SentencePiece encode returns a tuple, so leave it False for that path. Get it wrong and you'll see a KeyError: 'input_ids'.

Wiring it in

The output is a PTVDATASET, another pack-private type. It goes into the Pt Data Loader node to batch it, then into the trainer. The pack's own fine-tuning doc walks the whole chain: tokenizer encode → this node → data loader → Pto AdamW optimizer → loss node → Pt Train Transformer Classification Model. With DistilBERT on IMDB you should land around 92-93% validation accuracy.

One quirk worth knowing: the tokenizer you wire in should match the model you train, and for the HF tokenizer you can set max_length to 0 to use the model's own maximum.

Install

Same as every node in this pack - one install, all the nodes come along:

cd ComfyUI/custom_nodes
git clone https://github.com/HowToSD/ComfyUI-Pt-Wrapper

Restart ComfyUI. Or search "ComfyUI-Pt-Wrapper" in ComfyUI Manager. Requirements include transformers>=4.47.1 and datasets>=3.4.1, so expect a chunky first install.

Common issues

  • KeyError: 'input_ids' - encode_return_dict is False but your tokenizer returns a dict. Flip it.
  • First run downloads the dataset (IMDB is ~80MB) and, if you use an HF tokenizer, the tokenizer files too. It's a one-time thing.
  • Wrong encode wired in - the node does no checking; if you feed it a callable that doesn't return (tokens, masks), you'll get a confusing unpacking error. Wire the tokenizer node correctly.

This node is the right pick whenever your text data lives on the HF Hub. If your data is a local JSONL file instead, the pack has Ptv Hf Local Dataset for exactly that.

CategoryTraining

Inputs (7)

NameTypeDefaultDescription
nameSTRING
splitSTRINGtrain
sample_field_nameSTRINGtext
label_field_nameSTRINGlabel
encodePTCALLABLE
remove_html_tagsBOOLEANfalse
encode_return_dictBOOLEANfalse

Outputs (1)

NameTypeDescription
PTVDATASETPTVDATASET