Ptv Hf Dataset With Token Encode
Text classification data, straight from Hugging Face
- encode
- PTVDATASET
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. Usetest(or whatever the dataset calls it) for evaluation. - sample_field_name / label_field_name - defaults
textandlabel. If your dataset calls themreviewandsentiment, change these. - encode - the
PTCALLABLEfrom 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 aKeyError: '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_dictisFalsebut 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.
Inputs (7)
| Name | Type | Default | Description |
|---|---|---|---|
| name | STRING | — | |
| split | STRING | train | — |
| sample_field_name | STRING | text | — |
| label_field_name | STRING | label | — |
| encode | PTCALLABLE | — | |
| remove_html_tags | BOOLEAN | false | — |
| encode_return_dict | BOOLEAN | false | — |
Outputs (1)
| Name | Type | Description |
|---|---|---|
| PTVDATASET | PTVDATASET | — |