Ptn Hf Fine Tuned Classification Model

This is the node you reach for when you want to fine-tune a real Hugging Face transformer for text classification and you don't want to write a training loop. It drops a pretrained model like distilbert-base-uncased into the node graph, bolts a single-output classification head onto it, and hands you back a PTMODEL you can train with the pack's Pt Train Classification Transformer Model node. The author's own example hits 92–93% validation accuracy on the IMDB sentiment dataset with DistilBERT - that's not a toy.

How it works

Pick a model_name and the node loads it via AutoConfig + the matching from_pretrained call. Here's the catch worth knowing before you type anything: it only resolves the five model families hard-coded in the source - bert-base-uncased, roberta-base, distilbert-base-uncased, albert-base-v2, and google/electra-base-discriminator. Anything else throws a "model not supported" error. Weights are downloaded from Hugging Face on first run, so have a network connection and a few hundred MB of patience.

After the transformer, a nn.Linear(hidden_size, 1) head squeezes everything to one logit - it's a binary classifier, and it pairs with Ptn BCE With Logits Loss, which has the sigmoid built in. The head's weights are Xavier-initialized fresh, so you're only starting the classification layer from scratch; the transformer keeps its pretrained weights.

The inputs that matter

• model_name - the HF hub ID, and it must start with one of the five supported prefixes above. Use the exact same name in the Hf Tokenizer Encode node, because a BERT tokenizer won't produce IDs a RoBERTa model understands.
• use_mean_pooling (default True) - averages the hidden states over non-padding tokens using the attention mask. Turn it off and it falls back to the pooler output, then the CLS token. Mean pooling is the safer default for these encoder models.
• dropout (default 0) - applied before the final linear layer. 0 disables it; the example workflows mostly run without.

Everything comes out the single PTMODEL output, which wires into a chained model or straight into the trainer.

Wiring it up

The README's distilbert_classification.json is the canonical layout: Hf Tokenizer Encode and Ptv Hf Dataset With Token Encode load and tokenize the text, Pt Data Loader batches it, and Ptn Hf Fine Tuned Classification Model feeds the Pt Train Classification Transformer Model node alongside an Pto AdamW optimizer and the BCE-with-logits loss. If you want to freeze the transformer for a few warmup epochs first, swap in Pt Train Fine Tune Classification Transformer Model and set its freeze-epochs parameter.

Installing

Same story as every node in this pack - ComfyUI Manager → search "Pt-Wrapper" (pack title: ComfyUI-Pt-Wrapper), or:

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

Restart ComfyUI. The pack's requirements.txt is heavy - transformers, datasets, peft, accelerate, sentencepiece - so the first launch takes a while. Don't skip the "first run downloads the model" step; it's normal, not a hang.

Where people get burned

The biggest trap is a model name that isn't in the supported list - you'll get a clean ValueError, but only after downloading a config. Another one: running this node with use_mean_pooling off on a model without a pooler (DistilBERT has none) silently falls back to CLS-style [:, 0, :] extraction, which is fine but not what you expected. And remember every run re-instantiates the model - IS_CHANGED returns NaN, so the graph re-executes the node on every queue. Fine for training; wasteful if you only wanted to inspect it.