Nodes/ComfyUI-Pt-Wrapper/Ptn Hf Lora Classification Model
ComfyUI Node

Ptn Hf Lora Classification Model

LoRA fine-tuning for HF text classifiers, without touching code

By HowToSD·Created about a year ago·Updated about a year ago· 7
Ptn Hf Lora Classification Model
    • PTMODEL
    model_name
    use_mean_poolingtrue
    dropout0.000
    lora_r8
    lora_alpha16
    lora_dropout0.100

    This is PtnHfFineTunedClassificationModel's smarter sibling. Same idea - a Hugging Face encoder with a binary classification head, no Python - except the pretrained transformer gets adapted with LoRA instead of full fine-tuning. The base weights stay frozen, and only a handful of low-rank adapter matrices plus the new head actually train. If you've used LoRA for image models, the mental model transfers cleanly: same rank/alpha tradeoff, just applied to attention and feedforward layers in a text model.

    How it works

    Under the hood it loads the same five supported model families (bert-base-uncased, roberta-base, distilbert-base-uncased, albert-base-v2, google/electra-base-discriminator - anything else errors out), then runs it through peft.get_peft_model with target_modules="all-linear". The author's config also sets bias="lora_only", so biases in LoRA-injected layers are trainable too. The effective LoRA contribution is (lora_alpha / lora_r) * (A @ B) - that alpha-over-rank scaling is exactly the knob you already know from image LoRAs. The classification head is a fresh nn.Linear(hidden_size, 1) with Xavier init, and base model weights are frozen throughout.

    Because the base is frozen, training is faster and needs less VRAM than the full fine-tune variant, and you get most of the accuracy back - the tradeoff you'd expect from any LoRA setup.

    The inputs that matter

    The first three are shared with the non-LoRA node:

    • model_name - the HF hub ID, restricted to the supported prefixes.
    • use_mean_pooling (default True) - average token embeddings over the attention mask instead of using the pooler/CLS token.
    • dropout (default 0) - dropout before the final linear head.

    Then the LoRA controls:

    • lora_r (default 8) - rank of the low-rank matrices. Higher rank = more capacity, more trainable params. 8 is a fine starting point for a classifier.
    • lora_alpha (default 16) - the scaling factor; the effective contribution scales as alpha/r, so 16 over 8 is an effective 2x.
    • lora_dropout (default 0.1) - dropout on the LoRA layer inputs.

    All of it exits as one PTMODEL, ready for the same trainer nodes the non-LoRA version uses.

    Wiring it up

    The flow is identical to PtnHfFineTunedClassificationModel: tokenize with Hf Tokenizer Encode (same model name!), load the dataset with Ptv Hf Dataset With Token Encode, batch with Pt Data Loader, then feed this model plus an Pto AdamW optimizer and Ptn BCE With Logits Loss into Pt Train Classification Transformer Model. For a first run, leave rank at 8, keep alpha at 16, and let it train.

    Installing

    Same as the rest of the pack. ComfyUI Manager → search "Pt-Wrapper", or:

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

    Restart and let the heavy requirements.txt (transformers, peft, accelerate, datasets) install. First run downloads the base model weights from Hugging Face - expect that to take a minute.

    Where people get burned

    Two things. First, the model-name restriction applies here too, so don't copy a fancy model ID in and wonder why it fails. Second, LoRA's whole selling point - that you can swap adapters onto a fixed base - is only useful if you save the adapter, so make sure the workflow includes a save step or you'll retrain from scratch every time. And the same IS_CHANGED = NaN gotcha applies: this node rebuilds on every queue, which is what you want during training but a silent cost if you're just poking at the graph.

    CategoryTraining

    Inputs (6)

    NameTypeDefaultDescription
    model_nameSTRING
    use_mean_poolingBOOLEANtrue
    dropoutFLOAT0.0000–1
    lora_rINT81–1000000
    lora_alphaINT161–1000000
    lora_dropoutFLOAT0.1000–1

    Outputs (1)

    NameTypeDescription
    PTMODELPTMODEL