diff --git a/README.md b/README.md index 1f6b9fd..e222c22 100644 --- a/README.md +++ b/README.md @@ -1,9 +1,20 @@ ![Logo](docs/images/main_logo.png) -## Update — March 28, 2025 +## Update — March 31, 2025 -**Now with Transformer support** -Train NLP models directly inside ComfyUI. See: [`embedding_transformer_classification.json`](examples/workflows/embedding_transformer_classification.json). Achieves ~83% accuracy on IMDB sentiment with Hugging Face datasets—using only node workflows. +**Now with Hugging Face pretrained model support for text classification** + +You can use **ComfyUI-Pt-Wrapper** to fine-tune Hugging Face's pretrain models for text classification, such as sentiment analysis on a movie review dataset. + +You can choose from the following models: +* bert-base-uncased +* roberta-base +* distilbert-base-uncased +* albert-base-v2 + +Performance was not compromised in adding this feature — you should achieve around 92-93% validation accuracy on the IMDB dataset with the pretrained DistilBERT model. + +Check out [Fine-tuning Huggning Face models for classification](docs/fine_tuning_huggning_face_models_for_classification.md). --- diff --git a/docs/fine_tuning_huggning_face_models_for_classification.md b/docs/fine_tuning_huggning_face_models_for_classification.md new file mode 100644 index 0000000..99f315d --- /dev/null +++ b/docs/fine_tuning_huggning_face_models_for_classification.md @@ -0,0 +1,57 @@ +# Fine-tuning Huggning Face models for classification + +You can use **ComfyUI-Pt-Wrapper** to fine-tune Hugging Face's pretrain models for text classification, such as sentiment analysis on a movie review dataset. + +You can choose from the following models: +* bert-base-uncased +* roberta-base +* distilbert-base-uncased +* albert-base-v2 + +You can find an example workflow at `examples/workflows/distilbert_classification.json`. + +The flow consists of the following groups of nodes: +* Loading data +* Setting up the model, loss, optimizer, and scheduler +* Training + +## Loading Data + +This phase consists of two steps: +* Specify a Hugging Face tokenizer encode for converting text to token IDs +* Specify a Hugging Face text dataset to load + +## Specify a Hugging Face tokenizer encode for converting text to token IDs +Use **Hf Tokenizer Encode** to load the tokenizer for text to token ID convertion. +This node returns the encode function that dataset uses to the actual tokenization. +Note that each tokenizer is specific to the model. So you need to specify the same name as the model. You can specify one of the models as listed in the beginning of this document. + In the example workflow, a string node was created to specify the model name and this is connected to both the Encode node and the Model node. +Note that you can set max_length to 0 to use the model specific max length. + +## Specify a Hugging Face text dataset to load +Specify **Ptv Hf Dataset With Token Encode** to load the text data, and encode each word into token IDs using the Encode function above. **Select encode_return_dict to True.** + +Then use **Pt Data Loader** to batch the data and feed it into the training node. + +## Setting Up Optimizer and Scheduler + +This workflow uses **Pto AdamW** as the optimizer. You can also experiment with other optimizers and schedulers included in this extension. + +## Loss + +Use **Ptn BCE with Logits Loss**, which includes a built-in sigmoid activation, allowing you to pass raw logits directly from the model. + +## Model + +Use an Transformer model with a linear head (**Ptn Hf Fine Tuned Classification Model**). +This model consists of: +* Hugging Face pretrained model +* Linear head that outputs a logit for binary classification + +## Training + +Specify **Pt Train Transformer Classification Model** and configure training parameters. + +## Expected Accuracy + +Using this workflow, you should achieve around 92-93% validation accuracy on the Hugging Face version of the IMDB dataset with the pretrained DistilBERT model. \ No newline at end of file diff --git a/docs/model_training.md b/docs/model_training.md index 08bcf09..f6a5db5 100644 --- a/docs/model_training.md +++ b/docs/model_training.md @@ -110,4 +110,8 @@ Refer to `examples/workflows/train_model_with_custom_loss_example.json` for an e Refer to [Training RNN for text classification](training_rnn_for_classification.md) # Training Transformer -Refer to [Training Transformer for text classification](training_transformer_for_classification.md) \ No newline at end of file +Refer to [Training Transformer for text classification](training_transformer_for_classification.md) + +# Fine-tuning Huggning Face models for classification +Refer to [Fine-tuning Huggning Face models for classification] +(fine_tuning_huggning_face_models_for_classification.md). \ No newline at end of file