Added Hugging Face pretrained model support.

.gitignore

@@ -5,6 +5,7 @@ models/
 /datasets/
 examples/datasets/.~lock*
 .vscode
+output.txt
 
 # But don't ignore these paths:
 !models/sentencepiece/

__init__.py

@@ -1,5 +1,6 @@
 # This file is auto-generated. Do not edit manually.
 from typing import Dict, Type, TypeVar
+from .modules.hugging_face_wrapper.hf_tokenizer_encode import HfTokenizerEncode
 from .modules.pytorch_wrapper.pt_abs import PtAbs
 from .modules.pytorch_wrapper.pt_acos import PtAcos
 from .modules.pytorch_wrapper.pt_add import PtAdd
@@ -111,6 +112,7 @@ from .modules.pytorch_wrapper.pt_tokenizer import PtTokenizer
 from .modules.pytorch_wrapper.pt_train_classification_model import PtTrainClassificationModel
 from .modules.pytorch_wrapper.pt_train_classification_model_lr import PtTrainClassificationModelLr
 from .modules.pytorch_wrapper.pt_train_classification_transformer_model import PtTrainClassificationTransformerModel
+from .modules.pytorch_wrapper.pt_train_fine_tune_classification_transformer_model import PtTrainFineTuneClassificationTransformerModel
 from .modules.pytorch_wrapper.pt_train_model import PtTrainModel
 from .modules.pytorch_wrapper.pt_train_regression_model import PtTrainRegressionModel
 from .modules.pytorch_wrapper.pt_train_rnn_model import PtTrainRNNModel
@@ -163,6 +165,7 @@ from .modules.pytorch_wrapper.ptn_embedding_rnn_linear import PtnEmbeddingRNNLin
 from .modules.pytorch_wrapper.ptn_embedding_transformer_linear import PtnEmbeddingTransformerLinear
 from .modules.pytorch_wrapper.ptn_gru import PtnGRU
 from .modules.pytorch_wrapper.ptn_gru_linear import PtnGRULinear
+from .modules.pytorch_wrapper.ptn_hf_fine_tuned_classification_model import PtnHfFineTunedClassificationModel
 from .modules.pytorch_wrapper.ptn_huber_loss import PtnHuberLoss
 from .modules.pytorch_wrapper.ptn_instance_norm2d import PtnInstanceNorm2d
 from .modules.pytorch_wrapper.ptn_kl_div_loss import PtnKLDivLoss
@@ -209,6 +212,7 @@ NODE_CLASS_MAPPINGS (Dict[str, Type[T]]):
 """
 
 NODE_CLASS_MAPPINGS: Dict[str, Type[T]] = {
+    "HfTokenizerEncode": HfTokenizerEncode,
     "PtAbs": PtAbs,
     "PtAcos": PtAcos,
     "PtAdd": PtAdd,
@@ -320,6 +324,7 @@ NODE_CLASS_MAPPINGS: Dict[str, Type[T]] = {
     "PtTrainClassificationModel": PtTrainClassificationModel,
     "PtTrainClassificationModelLr": PtTrainClassificationModelLr,
     "PtTrainClassificationTransformerModel": PtTrainClassificationTransformerModel,
+    "PtTrainFineTuneClassificationTransformerModel": PtTrainFineTuneClassificationTransformerModel,
     "PtTrainModel": PtTrainModel,
     "PtTrainRNNModel": PtTrainRNNModel,
     "PtTrainRegressionModel": PtTrainRegressionModel,
@@ -372,6 +377,7 @@ NODE_CLASS_MAPPINGS: Dict[str, Type[T]] = {
     "PtnEmbeddingTransformerLinear": PtnEmbeddingTransformerLinear,
     "PtnGRU": PtnGRU,
     "PtnGRULinear": PtnGRULinear,
+    "PtnHfFineTunedClassificationModel": PtnHfFineTunedClassificationModel,
     "PtnHuberLoss": PtnHuberLoss,
     "PtnInstanceNorm2d": PtnInstanceNorm2d,
     "PtnKLDivLoss": PtnKLDivLoss,
@@ -418,6 +424,7 @@ NODE_DISPLAY_NAME_MAPPINGS (Dict[str, str]):
 """
 
 NODE_DISPLAY_NAME_MAPPINGS: Dict[str, str] = {
+    "HfTokenizerEncode": "Hf Tokenizer Encode",
     "PtAbs": "Pt Abs",
     "PtAcos": "Pt Acos",
     "PtAdd": "Pt Add",
@@ -529,6 +536,7 @@ NODE_DISPLAY_NAME_MAPPINGS: Dict[str, str] = {
     "PtTrainClassificationModel": "Pt Train Classification Model",
     "PtTrainClassificationModelLr": "Pt Train Classification Model Lr",
     "PtTrainClassificationTransformerModel": "Pt Train Classification Transformer Model",
+    "PtTrainFineTuneClassificationTransformerModel": "Pt Train Fine Tune Classification Transformer Model",
     "PtTrainModel": "Pt Train Model",
     "PtTrainRNNModel": "Pt Train RNN Model",
     "PtTrainRegressionModel": "Pt Train Regression Model",
@@ -581,6 +589,7 @@ NODE_DISPLAY_NAME_MAPPINGS: Dict[str, str] = {
     "PtnEmbeddingTransformerLinear": "Ptn Embedding Transformer Linear",
     "PtnGRU": "Ptn GRU",
     "PtnGRULinear": "Ptn GRU Linear",
+    "PtnHfFineTunedClassificationModel": "Ptn Hf Fine Tuned Classification Model",
     "PtnHuberLoss": "Ptn Huber Loss",
     "PtnInstanceNorm2d": "Ptn Instance Norm 2d",
     "PtnKLDivLoss": "Ptn KL Div Loss",

docs/reference/hf_tokenizer_encode.md

@@ -0,0 +1,35 @@
+# Hf Tokenizer Encode
+Hugging Face tokenizer wrapper for encoding text into token ID tensors.
+
+    Args:  
+        model_name (str): Name of the model.
+        padding (bool): Whether to pad sequences to equal length.
+        padding_method (str): 'longest' or 'max_length'.
+            - 'longest': pad to the longest sequence length after truncation in this batch.
+            - 'max_length': pad all sequences to the fixed length `max_length`.
+        truncation (bool): Whether to truncate sequences to `max_length`.
+        max_length (int): The target length for truncation or padding.  
+                          If you specify 0, it uses model's maximum input length.
+
+    Returns:  
+        PTCALLABLE: A tuple containing a callable that accepts a sentence or list of
+        sentences and returns token ID tensor(s) and attention mask(s).
+
+## Input
+| Name | Data type |
+|---|---|
+| model_name | String |
+| padding | Boolean |
+| padding_method |  |
+| truncation | Boolean |
+| max_length | Int |
+
+## Output
+| Data type |
+|---|
+| Ptcallable |
+
+<HR>
+Category: PyTorch wrapper - Training
+
+ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.

docs/reference/node_reference.md

@@ -160,6 +160,7 @@
 | [Ptn Embedding Transformer Linear](ptn_embedding_transformer_linear.md) | A Transformer encoder model with a linear head. |
 | [Ptn GRU](ptn_gru.md) | A gated recurrent unit (GRU) model consisting of one or more of a recurrent layer.   |
 | [Ptn GRU Linear](ptn_gru_linear.md) | A recurrent neural network (GRU) model with a linear head. |
+| [Ptn Hf Fine Tuned Classification Model](ptn_hf_fine_tuned_classification_model.md) | A binary classification model containing a Hugging Face pretrained Transformer model. |
 | [Ptn Instance Norm 2d](ptn_instance_norm2d.md) | A normalization model to normalize elements over spatial axes within each channel for each sample. |
 | [Ptn LSTM](ptn_lstm.md) | A long short-term memory  (LSTM) model consisting of one or more of a recurrent layer.   |
 | [Ptn LSTM Linear](ptn_lstm_linear.md) | A recurrent neural network (LSTM) model with a linear head. |
@@ -225,6 +226,7 @@
 ## PyTorch wrapper - Training
 | Node | Description |
 | --- | --- |
+| [Hf Tokenizer Encode](hf_tokenizer_encode.md) | Hugging Face tokenizer wrapper for encoding text into token ID tensors. |
 | [Pt Apply Function](pt_apply_function.md) | Applies a function to the input tensor. |
 | [Pt Compute Loss](pt_compute_loss.md) | Computes loss for the input with the target using the specified loss function. |
 | [Pt Data Loader](pt_data_loader.md) | Loads data from a dataset node and creates a PyTorch DataLoader.   |
@@ -236,6 +238,7 @@
 | [Pt Train Classification Model](pt_train_classification_model.md) | Trains a classification model using a given dataset, optimizer, and number of epochs.   |
 | [Pt Train Classification Model Lr](pt_train_classification_model_lr.md) | Trains a classification model using a given dataset, optimizer, and number of epochs with learning rate decay.   |
 | [Pt Train Classification Transformer Model](pt_train_classification_transformer_model.md) | Trains a classification Transformer model using a given dataset, loss function, optimizer, and number of epochs with learning rate decay.   |
+| [Pt Train Fine Tune Classification Transformer Model](pt_train_fine_tune_classification_transformer_model.md) | Fine-tunes a classification Transformer model using a given dataset, loss function, optimizer, and number of epochs with learning rate decay.   |
 | [Pt Train Model](pt_train_model.md) | Trains a model using a given dataset, loss function, optimizer, and number of epochs with learning rate decay.   |
 | [Pt Train RNN Model](pt_train_rnn_model.md) | Trains an RNN model using a given dataset, loss function, optimizer, and number of epochs with learning rate decay.   |
 | [Pt Train Regression Model](pt_train_regression_model.md) | Trains a regression model using a given dataset, optimizer, and number of epochs.   |

docs/reference/pt_train_fine_tune_classification_transformer_model.md

@@ -0,0 +1,45 @@
+# Pt Train Fine Tune Classification Transformer Model
+Fine-tunes a classification Transformer model using a given dataset, loss function, optimizer, and number of epochs with learning rate decay.  
+
+   Args:  
+        model (PTMODEL): The PyTorch model to be trained.  
+        train_loader (torch.utils.data.DataLoader): DataLoader containing the training dataset.  
+        optimizer (PTOPTIMIZER): Optimizer used for training.  Only Adam and AdamW are supported.
+        loss (PTLOSS): The loss class.  
+        epochs (int): Number of training epochs.  
+        freeze_pretrained_module_epochs (int) : Number of epochs to freeze pretrained module.
+        use_gpu (bool): Whether to use GPU for training.  
+        early_stopping (bool): Whether to turn on early stopping  
+        early_stopping_rounds (int): Specifies the number of rounds to monitor the evaluation loss when early_stopping is enabled and validation data is specified. If the loss does not decrease compared to the current best loss within this period, training will be stopped.  
+        output_best_val_model (bool): Whether to outputs the model with the best val loss or the model from the last epoch.  
+        classification_metrics (bool): If True and if val_loader is set, print out classification metrics.  This is only valid for classification model training.  
+        val_loader (PTDATALOADER): DataLoader for validation dataset (optional).
+
+## Input
+| Name | Data type |
+|---|---|
+| model | Ptmodel |
+| train_loader | Ptdataloader |
+| optimizer | Ptoptimizer |
+| loss_function | Ptloss |
+| epochs | Int |
+| freeze_pretrained_module_epochs | Int |
+| use_gpu | Boolean |
+| early_stopping | Boolean |
+| early_stopping_rounds | Int |
+| output_best_val_model | Boolean |
+| classification_metrics | Boolean |
+| scheduler | Ptlrscheduler |
+| val_loader | Ptdataloader |
+
+## Output
+| Data type |
+|---|
+| Ptmodel |
+| Tensor |
+| Tensor |
+
+<HR>
+Category: PyTorch wrapper - Training
+
+ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.

docs/reference/ptn_hf_fine_tuned_classification_model.md

@@ -0,0 +1,26 @@
+# Ptn Hf Fine Tuned Classification Model
+A binary classification model containing a Hugging Face pretrained Transformer model.
+
+        Args:
+        model_name (str): Name or path of the pretrained Hugging Face model.
+        use_mean_pooling (bool): If True, average token embeddings using the attention mask.
+        Otherwise, use pooler output or CLS token embedding.
+        dropout (float): Dropout rate applied before the final linear layer.
+                         Set to 0.0 to disable dropout. Must be ≥ 0.
+
+## Input
+| Name | Data type |
+|---|---|
+| model_name | String |
+| use_mean_pooling | Boolean |
+| dropout | Float |
+
+## Output
+| Data type |
+|---|
+| Ptmodel |
+
+<HR>
+Category: PyTorch wrapper - Model
+
+ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.

docs/reference/ptv_hf_dataset_with_token_encode.md

@@ -8,6 +8,7 @@ A PyTorch Dataset class wrapper for a Hugging Face dataset that converts text to
         label_field_name (str): Field name for labels.
         encode (PTCALLABLE): The reference to a token encoder function.
         remove_html_tags (bool): Remove html tags in text if True.
+        encode_return_dict (bool): Encode function returns a Dict instead of a Tuple.
 
 ## Input
 | Name | Data type |
@@ -18,6 +19,7 @@ A PyTorch Dataset class wrapper for a Hugging Face dataset that converts text to
 | label_field_name | String |
 | encode | Ptcallable |
 | remove_html_tags | Boolean |
+| encode_return_dict | Boolean |
 
 ## Output
 | Data type |

docs/training_rnn_for_classification.md

@@ -46,7 +46,7 @@ You can also enable `classification_metrics` to print validation accuracy to the
 
 ## Expected Accuracy
 
-Using this workflow, you should get around **83% validation accuracy** on the Hugging Face version of the IMDB dataset, which is a solid result for an RNN.
+Using this workflow, you should get around **85% validation accuracy** on the Hugging Face version of the IMDB dataset, which is a solid result for an RNN.
 Both GRU & LSTM workflows achieve approximately **87% validation accuracy**.
 
 # RNN Workflow with Training Embedding from Scratch

modules/hugging_face_wrapper/hf_dataset_with_token_encode.py

@@ -0,0 +1,89 @@
+from typing import Tuple, Callable, Optional
+import re
+from html import unescape
+import torch
+from torch.utils.data import Dataset
+from datasets import load_dataset  # Hugging Face (not PyTorch)
+
+
+class HfDatasetWithTokenEncode(Dataset):
+    """
+    Hugging Face Dataset wrapper with custom token encode function.
+
+    pragma: skip_doc
+    """
+    def __init__(self,
+                 dataset_name: str,
+                 split: str,
+                 sample_field_name: str = "text",
+                 label_field_name: str = "label",
+                 encode: Optional[Callable]=None,
+                 remove_html_tags:Optional[bool]=False,
+                 encode_return_dict:Optional[bool]=False) -> None:
+        """
+        Initializes the dataset by loading Hugging Face dataset.
+
+        Args:
+            dataset_name (str): The name of the Hugging Face dataset.
+            split (str): The dataset split (e.g., 'train').
+             sample_field_name (str): Field name for text samples.
+            label_field_name (str): Field name for labels.
+            encode (Optional[Callable]): Token encode function
+            remove_html_tags (Optional[bool]): Remove html tags in text if True.
+            encode_return_dict (Optional[bool]): Encode function returns a Dict instead of a Tuple.
+        """
+        self.dataset = load_dataset(dataset_name, split=split)
+        self.sample_field_name = sample_field_name
+        self.label_field_name = label_field_name
+        self.encode = encode
+        self.remove_html_tags = remove_html_tags
+        self.encode_return_dict = encode_return_dict
+
+    def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Retrieves the (tokens, label) pair at the given index.
+
+        Args:
+            idx (int): Index of the item.
+
+        Returns:
+            Tuple[Tuple[torch.Tensor, torch.Tensor], torch.Tensor]: Tuple of the following:  
+                - Tokens tensor and masks tensor.
+                - Label tensor
+        """
+        # Get a text sample from database and convert to a word embedding
+        sample = self.dataset[idx]
+        if self.remove_html_tags:
+            # Decode HTML entities like "&", "<", etc. to their actual characters
+            text = unescape(sample[self.sample_field_name])
+            
+            # Remove HTML tags (like <br />, <p>, etc.)
+            # The regex "<.*?>" uses .*? (non-greedy match) to ensure it matches one tag at a time,
+            # instead of greedily removing everything between the first '<' and last '>'
+            text = re.sub(r"<.*?>", " ", text)
+
+            # Replace consecutive white space to a single white space
+            text = re.sub(r"  +", " ", text)
+        else:
+            text = sample[self.sample_field_name]
+
+        if self.encode_return_dict:
+            output = self.encode(text)
+            tokens = output["input_ids"]
+            masks = output["attention_mask"]
+        else:
+            tokens, masks = self.encode(text)
+
+        # Get a corresponding label
+        label = torch.tensor(sample[self.label_field_name], dtype=torch.long)
+        return (tokens.squeeze(0), masks.squeeze(0)), label.squeeze(0)
+
+
+    def __len__(self) -> int:
+        """
+        Returns the number of samples in the dataset.
+
+        Returns:
+            int: Dataset length.
+        """
+        return len(self.dataset)

modules/hugging_face_wrapper/hf_glove_dataset.py

@@ -0,0 +1,132 @@
+from typing import Tuple
+import re
+from html import unescape
+import torch
+from torch.utils.data import Dataset
+from datasets import load_dataset  # Hugging Face (not PyTorch)
+import scipy
+import numpy as np
+# Monkey-patch gensim as it relies on scipy's deprecated triu method
+# /gensim/matutils.py has this line which causes ImportError:
+# from scipy.linalg import get_blas_funcs, triu
+# TODO: Monitor gensim and fix this if upgrade is made
+scipy.linalg.triu = np.triu
+import gensim.downloader as api
+
+class HfGloveDataset(Dataset):
+    """
+    Hugging Face Dataset wrapper with word embedding support.
+
+    pragma: skip_doc
+    """
+    def __init__(self,
+                 dataset_name: str,
+                 split: str,
+                 glove_dim: int = 100,
+                 max_seq_len: int = 256,
+                 sample_field_name: str = "text",
+                 label_field_name: str = "label") -> None:
+        """
+        Initializes the dataset by loading Hugging Face dataset and GloVe vectors.
+
+        Args:
+            dataset_name (str): The name of the Hugging Face dataset.
+            split (str): The dataset split (e.g., 'train').
+            glove_dim (int): Dimensionality of GloVe embeddings.
+            max_seq_len (int): Maximum sequence length.
+            sample_field_name (str): Field name for text samples.
+            label_field_name (str): Field name for labels.
+        """
+        self.dataset = load_dataset(dataset_name, split=split)
+        self.glove_dim = glove_dim
+        self.max_seq_len = max_seq_len
+        self.sample_field_name = sample_field_name
+        self.label_field_name = label_field_name
+
+        # Load GloVe from gensim
+        glove_name = f"glove-wiki-gigaword-{glove_dim}"
+        self.glove = api.load(glove_name)
+        self.unk_vector = torch.zeros(glove_dim)
+
+    def tokenize(self, text: str) -> list[str]:
+        """
+        Tokenizes raw text into a list of lowercase word tokens.
+        It removes HTML tags, decodes HTML entities, strips punctuation 
+        (while keeping letters and numbers), and splits on whitespace.
+
+        Args:
+            text (str): Raw text input (may include HTML and punctuation).
+
+        Returns:
+            list[str]: List of clean, lowercase word tokens.
+        """
+        # Decode HTML entities like "&", "<", etc. to their actual characters
+        text = unescape(text)
+        
+        # Remove HTML tags (like <br />, <p>, etc.)
+        # The regex "<.*?>" uses .*? (non-greedy match) to ensure it matches one tag at a time,
+        # instead of greedily removing everything between the first '<' and last '>'
+        text = re.sub(r"<.*?>", " ", text)
+        
+        # Remove all characters except letters, digits, and whitespace
+        # This strips punctuation like "!", ".", "(", but keeps words and numbers
+        text = re.sub(r"[^a-zA-Z0-9\s]", " ", text)
+
+        # Replace consecutive white space to a single white space
+        text = re.sub(r"  +", " ", text)
+
+        # Convert to lowercase and split by whitespace to get token list
+        return text.lower().split()
+
+
+    def vectorize(self, tokens: list[str]) -> torch.Tensor:
+        """
+        Converts tokens to a fixed-length tensor of word vectors.
+
+        Args:
+            tokens (list[str]): A list of word tokens.
+
+        Returns:
+            torch.Tensor: A tensor of shape (max_seq_len, glove_dim).
+        """
+        vectors = []
+        for token in tokens[:self.max_seq_len]:
+            if token in self.glove:
+                vec = torch.tensor(self.glove[token], dtype=torch.float32)
+            else:
+                vec = self.unk_vector
+            vectors.append(vec)
+
+        if len(vectors) < self.max_seq_len:
+            pad = [torch.zeros(self.glove_dim) for _ in range(self.max_seq_len - len(vectors))]
+            vectors.extend(pad)
+
+        return torch.stack(vectors)
+
+    def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Retrieves the (embedded sequence, label) pair at the given index.
+
+        Args:
+            idx (int): Index of the item.
+
+        Returns:
+            Tuple[torch.Tensor, torch.Tensor]: Embedding tensor and label tensor.
+        """
+        # Get a text sample from database and convert to a word embedding
+        sample = self.dataset[idx]
+        tokens = self.tokenize(sample[self.sample_field_name])
+        emb_tensor = self.vectorize(tokens)
+
+        # Get a corresponding label
+        label = torch.tensor(sample[self.label_field_name], dtype=torch.long)
+        return emb_tensor, label
+
+    def __len__(self) -> int:
+        """
+        Returns the number of samples in the dataset.
+
+        Returns:
+            int: Dataset length.
+        """
+        return len(self.dataset)

modules/hugging_face_wrapper/hf_tokenizer_encode.py

@@ -0,0 +1,134 @@
+from typing import Any, Dict, Union, List, Tuple, Callable
+import torch
+from transformers import AutoTokenizer
+
+
+class HfTokenizerEncode:
+    """
+    Hugging Face tokenizer wrapper for encoding text into token ID tensors.
+
+        Args:  
+            model_name (str): Name of the model.
+            padding (bool): Whether to pad sequences to equal length.
+            padding_method (str): 'longest' or 'max_length'.
+                - 'longest': pad to the longest sequence length after truncation in this batch.
+                - 'max_length': pad all sequences to the fixed length `max_length`.
+            truncation (bool): Whether to truncate sequences to `max_length`.
+            max_length (int): The target length for truncation or padding.  
+                              If you specify 0, it uses model's maximum input length.
+
+        Returns:  
+            PTCALLABLE: A tuple containing a callable that accepts a sentence or list of
+            sentences and returns token ID tensor(s) and attention mask(s).  
+
+    category: PyTorch wrapper - Training
+    """
+
+    @classmethod
+    def INPUT_TYPES(cls) -> Dict[str, Any]:
+        """
+        Defines the input types for the function.
+
+        Returns:
+            Dict[str, Any]: A dictionary specifying required input types.
+        """
+        return {
+            "required": {
+                "model_name": ("STRING", {"default":"", "multiline":False}),
+                "padding": ("BOOLEAN", {"default": True}),
+                "padding_method": (("max_length", "longest"),),
+                "truncation": ("BOOLEAN", {"default": True}),
+                "max_length": ("INT", {"default": 512, "min": 0, "max": 1_000_000}),
+            }
+        }
+
+    RETURN_TYPES: Tuple[str] = ("PTCALLABLE",)
+    FUNCTION: str = "f"
+    CATEGORY: str = "Data Analysis"
+
+    def encode(
+        self,
+        sentence: Union[str, List[str]],
+    ) -> Dict[str, torch.Tensor]:
+
+        """
+        Encodes a string or list of strings into token ID tensors.
+
+        Args:
+            sentence (Union[str, List[str]]): Input sentence(s).
+
+        Returns:
+            Dict:
+             input_ids: token IDs
+             attention_mask: Masks. 1 is non-padding, 0 is padding.
+        """
+        if self.padding:
+            if self.max_length == 0:
+                retval = self.tokenizer(
+                            sentence,
+                            padding=self.padding_method,
+                            truncation=self.truncation,
+                            return_tensors="pt"
+                        )
+            else:
+                retval = self.tokenizer(
+                            sentence,
+                            padding=self.padding_method,
+                            truncation=self.truncation,
+                            max_length=self.max_length,
+                            return_tensors="pt"
+                        )
+        else:  # no padding
+            if self.max_length == 0:
+                retval = self.tokenizer(
+                            sentence,
+                            padding=False,
+                            truncation=self.truncation,
+                            return_tensors="pt"
+                        )
+            else:
+                retval = self.tokenizer(
+                            sentence,
+                            padding=False,
+                            truncation=self.truncation,
+                            return_tensors="pt"
+                        )
+        return retval
+
+    def f(
+        self,
+        model_name: str,
+        padding: bool,
+        padding_method: str,
+        truncation: bool,
+        max_length: int
+    ) -> Tuple[Callable[[Union[str, List[str]]], Tuple[Union[torch.Tensor, List[torch.Tensor]], Union[torch.Tensor, List[torch.Tensor]]]]]:
+        """
+        Instantiates a Hugging Face tokenizer, binds configuration and returns a callable encoder function.
+        - Supports includes:
+        - BertTokenizer
+        - RobertaTokenizer
+        - DistilBertTokenizer
+        - AlbertTokenizer
+        - ElectraTokenizer
+
+        Args:
+            model_name (str): Name of the model.
+            padding (bool): Whether to pad sequences to equal length.
+            padding_method (str): 'longest' or 'max_length'.
+                - 'longest': pad to the longest sequence length after truncation in this batch.
+                - 'max_length': pad all sequences to the fixed length `max_length`.
+            truncation (bool): Whether to truncate sequences to `max_length`.
+            max_length (int): The target length for truncation or padding.
+                              If you specify 0, it uses model's maximum input length.
+
+        Returns:
+            Tuple[Callable]: A tuple containing a callable for the tokenizer function.
+        """
+
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+        self.padding = padding
+        self.padding_method = padding_method
+        self.truncation = truncation
+        self.max_length = max_length
+        return (self.encode,)

modules/pytorch_wrapper/pt_train_classification_transformer_model.py

@@ -114,9 +114,6 @@ class PtTrainClassificationTransformerModel:
         best_val_loss = None
         rounds_since_best_val_loss = 0
 
-        all_preds = []
-        all_labels = []
-
         with torch.inference_mode(False):
             train_loss = []
             val_loss = []
@@ -157,6 +154,8 @@ class PtTrainClassificationTransformerModel:
                 
                 if val_loader:
                     with torch.no_grad():
+                        all_preds = []
+                        all_labels = []
                         model.train(False)  # TODO: Change to model.*e*v*a*l() once Comfy's security checker is fixed.
                         total_loss = 0
                         total_samples = len(val_loader.dataset)
@@ -180,7 +179,7 @@ class PtTrainClassificationTransformerModel:
 
                             if classification_metrics and y.squeeze().dim() == 1 :
                                 # Accuracy collection for binary classification
-                                preds = (y_hat > 0.5).int().detach().cpu().numpy()
+                                preds = (F.sigmoid(y_hat) > 0.5).int().detach().cpu().numpy()
                                 labels = y.int().detach().cpu().numpy()
                                 all_preds.extend(preds)
                                 all_labels.extend(labels)

modules/pytorch_wrapper/pt_train_fine_tune_classification_transformer_model.py

@@ -0,0 +1,253 @@
+import os
+import sys
+from typing import Any, Dict, Optional, Tuple, Union
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from sklearn.metrics import accuracy_score
+from .utils import set_requires_grad
+
+class PtTrainFineTuneClassificationTransformerModel:
+    """
+    Fine-tunes a classification Transformer model using a given dataset, loss function, optimizer, and number of epochs with learning rate decay.  
+
+       Args:  
+            model (PTMODEL): The PyTorch model to be trained.  
+            train_loader (torch.utils.data.DataLoader): DataLoader containing the training dataset.  
+            optimizer (PTOPTIMIZER): Optimizer used for training.  Only Adam and AdamW are supported.
+            loss (PTLOSS): The loss class.  
+            epochs (int): Number of training epochs.  
+            freeze_pretrained_module_epochs (int) : Number of epochs to freeze pretrained module.
+            use_gpu (bool): Whether to use GPU for training.  
+            early_stopping (bool): Whether to turn on early stopping  
+            early_stopping_rounds (int): Specifies the number of rounds to monitor the evaluation loss when early_stopping is enabled and validation data is specified. If the loss does not decrease compared to the current best loss within this period, training will be stopped.  
+            output_best_val_model (bool): Whether to outputs the model with the best val loss or the model from the last epoch.  
+            classification_metrics (bool): If True and if val_loader is set, print out classification metrics.  This is only valid for classification model training.  
+            val_loader (PTDATALOADER): DataLoader for validation dataset (optional).  
+
+    category: PyTorch wrapper - Training
+    """
+
+    @classmethod
+    def INPUT_TYPES(cls) -> Dict[str, Any]:
+        """
+        Defines the input types required for training the model.
+
+        Returns:
+            Dict[str, Any]: A dictionary specifying required and optional input types.
+        """
+        return {
+            "required": {
+                "model": ("PTMODEL", {}),
+                "train_loader": ("PTDATALOADER", {}),
+                "optimizer": ("PTOPTIMIZER", {}),
+                "loss_function": ("PTLOSS", {}),
+                "epochs": ("INT", {"default":2, "min":1, "max":1e6}),
+                "freeze_pretrained_module_epochs": ("INT", {"default":1, "min":1, "max":1e6}),
+                "use_gpu": ("BOOLEAN", {"default":False}),
+                "early_stopping": ("BOOLEAN", {"default":False}),
+                "early_stopping_rounds": ("INT", {"default":10, "min":1, "max":1000}),
+                "output_best_val_model": ("BOOLEAN", {"default":True}),
+                "classification_metrics": ("BOOLEAN", {"default":True})
+            },
+            "optional": {
+                "scheduler": ("PTLRSCHEDULER", {}),
+                "val_loader": ("PTDATALOADER", {})
+            }
+        }
+
+    RETURN_NAMES: tuple = ("Model", "train loss", "val loss")
+    RETURN_TYPES: tuple = ("PTMODEL","TENSOR","TENSOR")
+    FUNCTION: str = "f"
+    CATEGORY: str = "Training"
+
+    @classmethod
+    def IS_CHANGED(cls, **kw):
+        return float("NaN")
+
+    def f(self, 
+          model: nn.Module, 
+          train_loader: torch.utils.data.DataLoader, 
+          optimizer: torch.optim.Optimizer, 
+          loss_function: nn.Module,
+          epochs: int,
+          freeze_pretrained_module_epochs,
+          use_gpu: bool, 
+          early_stopping: bool,
+          early_stopping_rounds: int,
+          output_best_val_model:bool,
+          classification_metrics:bool,
+          val_loader: Optional[torch.utils.data.DataLoader] = None,
+          ) -> Tuple[nn.Module, torch.Tensor, torch.Tensor]:
+        """
+        Trains an Transformer model.
+
+        Args:
+            model (torch.nn.Module): The PyTorch model to be trained.
+            train_loader (torch.utils.data.DataLoader): DataLoader containing the training dataset.
+            optimizer (PTOPTIMIZER): Optimizer used for training.  Only Adam and AdamW are supported.
+            loss (torch.nn.Module): The loss class.
+            epochs (int): Number of training epochs.
+            freeze_pretrained_module_epochs (int) : Number of epochs to freeze pretrained module.
+            use_gpu (bool): Whether to use GPU for training.
+            early_stopping (bool): Whether to turn on early stopping
+            early_stopping_rounds (int): Specifies the number of rounds to monitor the evaluation loss when early_stopping is enabled and validation data is specified. If the loss does not decrease compared to the current best loss within this period, training will be stopped.
+            output_best_val_model (bool): Whether to outputs the model with the best val loss or the model from the last epoch.
+            classification_metrics (bool): If True and if val_loader is set, print out classification metrics.  This is only valid for classification model training.
+            val_loader (Optional[torch.utils.data.DataLoader]): DataLoader for validation dataset (optional).
+
+        Returns:
+            Tuple[torch.nn.Module, torch.Tensor, torch.Tensor]: A tuple containing the trained model, 
+            training loss history, and validation loss history.
+        """
+        checkpoint_dir = os.path.realpath(
+            os.path.join(os.path.dirname(__file__), "..", "..", "models", "tmp"))
+        print(f"Checkpoint tmp directory is set to {checkpoint_dir}")
+        if os.path.exists(checkpoint_dir) is False:
+            os.makedirs(checkpoint_dir)
+            print(f"Created {checkpoint_dir}")
+        best_checkpoint_path = os.path.join(checkpoint_dir, "best_model.pt")
+
+        best_val_loss = None
+        rounds_since_best_val_loss = 0
+
+        with torch.inference_mode(False):
+            train_loss = []
+            val_loss = []
+            if use_gpu:
+                model.to("cuda")
+
+            # Train the model
+            for epoch in range(epochs):
+                model.train()
+                total_loss = 0
+                total_samples = len(train_loader.dataset)
+
+                if epoch < freeze_pretrained_module_epochs:
+                    set_requires_grad(model.llm_model, False)  # freeze encoder
+                elif epoch == freeze_pretrained_module_epochs:
+                    set_requires_grad(model.llm_model, True)   # unfreeze encoder
+
+                    # Rebuild optimizer with separate LR for pretrained encoder and classifier head
+                    if isinstance(optimizer, torch.optim.Adam):
+                        lr = optimizer.defaults.get("lr", 1e-3)
+                        betas = optimizer.defaults.get("betas", (0.9, 0.999))
+                        eps = optimizer.defaults.get("eps", 1e-8)
+                        weight_decay = optimizer.defaults.get("weight_decay", 0)
+
+                        optimizer = torch.optim.Adam([
+                            {"params": model.llm_model.parameters(), "lr": lr * 0.2},
+                            {"params": model.linear.parameters(), "lr": lr}
+                        ], lr=lr, betas=betas, eps=eps, weight_decay=weight_decay)
+
+                    elif isinstance(optimizer, torch.optim.AdamW):
+                        lr = optimizer.defaults.get("lr", 1e-3)
+                        betas = optimizer.defaults.get("betas", (0.9, 0.999))
+                        eps = optimizer.defaults.get("eps", 1e-8)
+                        weight_decay = optimizer.defaults.get("weight_decay", 0)
+
+                        optimizer = torch.optim.AdamW([
+                            {"params": model.llm_model.parameters(), "lr": lr * 0.2},
+                            {"params": model.linear.parameters(), "lr": lr}
+                        ], lr=lr, betas=betas, eps=eps, weight_decay=weight_decay)
+
+                    else:
+                        raise ValueError("Only Adam and AdamW are supported for optimizer re-instantiation.")
+
+                for x, y in train_loader:
+                    if use_gpu:
+                        if isinstance(x, list):
+                            x = [e.to("cuda") for e in x]
+                        else:
+                            x = x.to("cuda")
+                        y = y.to("cuda")
+                    optimizer.zero_grad()
+
+                    y_hat = model(x)
+                    if y.dim() == 1:
+                        y = torch.unsqueeze(y, dim=-1)
+                        y = y.to(torch.float32)
+
+                    loss = loss_function(y_hat, y)
+                    loss.backward()
+                    optimizer.step()
+ 
+                    batch_size = y.size(0)
+                    batch_loss = loss.item() * batch_size
+                    total_loss += batch_loss
+
+                epoch_loss = total_loss / total_samples
+                print(f"Epoch (train) {epoch+1}/{epochs}, Loss: {epoch_loss:.4f}")
+                train_loss.append(epoch_loss)
+                
+                if val_loader:
+                    with torch.no_grad():
+                        all_preds = []
+                        all_labels = []
+
+                        model.train(False)  # TODO: Change to model.*e*v*a*l() once Comfy's security checker is fixed.
+                        total_loss = 0
+                        total_samples = len(val_loader.dataset)
+                        for x, y in val_loader:
+                            if use_gpu:
+                                if isinstance(x, list):
+                                    x = [e.to("cuda") for e in x]
+                                else:
+                                    x = x.to("cuda")
+                                y = y.to("cuda")
+
+                            y_hat = model(x)
+                            if y.dim() == 1:
+                                y = torch.unsqueeze(y, dim=-1)
+                                y = y.to(torch.float32)
+
+                            loss = loss_function(y_hat, y)
+                            batch_size = y.size(0)
+                            batch_loss = loss.item() * batch_size
+                            total_loss += batch_loss
+
+                            if classification_metrics and y.squeeze().dim() == 1 :
+                                # Accuracy collection for binary classification
+                                preds = (F.sigmoid(y_hat) > 0.5).int().detach().cpu().numpy()
+                                labels = y.int().detach().cpu().numpy()
+                                all_preds.extend(preds)
+                                all_labels.extend(labels)
+
+                        epoch_loss = total_loss / total_samples
+                        print(f"Epoch (val) {epoch+1}/{epochs}, Loss: {epoch_loss:.4f}")
+                        val_loss.append(epoch_loss)
+
+                        if classification_metrics:
+                            if len(all_labels):
+                                # Compute accuracy
+                                acc = accuracy_score(all_labels, all_preds)
+                                print(f"Validation Accuracy: {acc:.4f}")
+                            else:
+                                print("Accuracy was not computed as the label is not scalar.")
+
+                        if best_val_loss is None or epoch_loss < best_val_loss:
+                            best_val_loss = epoch_loss
+                            rounds_since_best_val_loss = 0
+                            torch.save(model.state_dict(), best_checkpoint_path)
+                            print(f"Updated best val loss. Saved the model as {best_checkpoint_path}")
+                        else:
+                            rounds_since_best_val_loss += 1
+                            if early_stopping:
+                                if rounds_since_best_val_loss >= early_stopping_rounds:
+                                    print(f"Val loss did not improve in {early_stopping_rounds} rounds. Stopping training.")
+                                    break
+
+            if use_gpu:
+                model.to("cpu")
+
+            if val_loader:
+                if epoch_loss > best_val_loss and output_best_val_model:
+                    sd = torch.load(best_checkpoint_path, weights_only=True)
+                    model.load_state_dict(sd)
+                    print("Loaded the best model for output")
+                else:
+                    print("Using the final epoch model for output.")
+
+            return (model, 
+                    torch.tensor(train_loss, dtype=torch.float32), 
+                    torch.tensor(val_loss, dtype=torch.float32))

modules/pytorch_wrapper/pt_train_rnn_model.py

@@ -127,9 +127,6 @@ class PtTrainRNNModel:
         best_val_loss = None
         rounds_since_best_val_loss = 0
 
-        all_preds = []
-        all_labels = []
-
         with torch.inference_mode(False):
             train_loss = []
             val_loss = []
@@ -190,6 +187,8 @@ class PtTrainRNNModel:
                 
                 if val_loader:
                     with torch.no_grad():
+                        all_preds = []
+                        all_labels = []
                         model.train(False)  # TODO: Change to model.*e*v*a*l() once Comfy's security checker is fixed.
                         total_loss = 0
                         total_samples = len(val_loader.dataset)
@@ -228,7 +227,7 @@ class PtTrainRNNModel:
 
                             if classification_metrics and y.squeeze().dim() == 1 :
                                 # Accuracy collection for binary classification
-                                preds = (y_hat > 0.5).int().detach().cpu().numpy()
+                                preds = (F.sigmoid(y_hat) > 0.5).int().detach().cpu().numpy()
                                 labels = y.int().detach().cpu().numpy()
                                 all_preds.extend(preds)
                                 all_labels.extend(labels)

modules/pytorch_wrapper/ptn_gru.py

@@ -35,7 +35,7 @@ class PtnGRU:
                 "num_layers": ("INT", {"default": 1, "min": 1, "max": 1e3}),
                 "bias": ("BOOLEAN", {"default": True}),
                 "batch_first": ("BOOLEAN", {"default": True}),
-                "dropout": ("FLOAT", {"default": 0, "min": 0, "max": 1}),
+                "dropout": ("FLOAT", {"default": 0, "min": 0, "max": 1, "step": 0.001}),
                 "bidirectional": ("BOOLEAN", {"default": False}),
             }
         }

modules/pytorch_wrapper/ptn_gru_linear.py

@@ -105,7 +105,7 @@ class PtnGRULinear:
                 "num_layers": ("INT", {"default": 1, "min": 1, "max": 1e3}),
                 "bias": ("BOOLEAN", {"default": True}),
                 "batch_first": ("BOOLEAN", {"default": True}),
-                "dropout": ("FLOAT", {"default": 0, "min": 0, "max": 1}),
+                "dropout": ("FLOAT", {"default": 0, "min": 0, "max": 1, "step": 0.001}),
                 "bidirectional": ("BOOLEAN", {"default": False}),
                 "linear_output_size": ("INT", {"default": 1, "min": 1, "max": 1e6}),
                 "linear_bias": ("BOOLEAN", {"default": True}),

modules/pytorch_wrapper/ptn_hf_fine_tuned_classification_model.py

@@ -0,0 +1,70 @@
+from typing import Any, Dict, Tuple
+import torch
+import torch.nn as nn
+
+from .ptn_hf_fine_tuned_classification_model_def import HfFineTunedClassificationModel
+
+class PtnHfFineTunedClassificationModel:
+    """
+    A binary classification model containing a Hugging Face pretrained Transformer model.
+
+            Args:
+            model_name (str): Name or path of the pretrained Hugging Face model.
+            use_mean_pooling (bool): If True, average token embeddings using the attention mask.
+            Otherwise, use pooler output or CLS token embedding.
+            dropout (float): Dropout rate applied before the final linear layer.
+                             Set to 0.0 to disable dropout. Must be ≥ 0. 
+
+    category: PyTorch wrapper - Model
+    """
+
+    @classmethod
+    def INPUT_TYPES(cls) -> Dict[str, Any]:
+        """
+        Defines the input types for the function.
+
+        Returns:
+            Dict[str, Any]: A dictionary specifying required input types.
+        """
+        return {
+            "required": {
+                "model_name": ("STRING", {"default": "", "multiline": False}),
+                "use_mean_pooling": ("BOOLEAN", {"default": True}),
+                "dropout": ("FLOAT", {"default": 0, "min": 0, "max": 1, "step": 0.001})
+            }
+        }
+
+    RETURN_TYPES: Tuple[str] = ("PTMODEL",)
+    FUNCTION: str = "f"
+    CATEGORY: str = "Training"
+
+    @classmethod
+    def IS_CHANGED(cls, **kw):
+        return float("NaN")
+
+    def f(self,
+          model_name: str,
+          use_mean_pooling: bool,
+          dropout: float
+        ) -> Tuple[nn.Module]:
+        """
+        Instantiates a binary classification model containing a Hugging Face pretrained Transformer model.
+
+            Args:
+            model_name (str): Name or path of the pretrained Hugging Face model.
+            use_mean_pooling (bool): If True, average token embeddings using the attention mask.
+            Otherwise, use pooler output or CLS token embedding.
+            dropout (float): Dropout rate applied before the final linear layer.
+                             Set to 0.0 to disable dropout. Must be ≥ 0.    
+                                
+            Returns:
+                Tuple[nn.Module]: A tuple containing the instantiated PyTorch model.
+        """
+        with torch.inference_mode(False):
+            model = HfFineTunedClassificationModel(
+                    model_name=model_name,
+                    use_mean_pooling=use_mean_pooling,
+                    dropout=dropout
+                )
+
+        return (model,)

modules/pytorch_wrapper/ptn_hf_fine_tuned_classification_model_def.py

@@ -0,0 +1,92 @@
+from typing import List, Dict, Any, Tuple, Type
+import torch
+import torch.nn as nn
+from transformers import AutoConfig, PreTrainedModel
+
+class HfFineTunedClassificationModel(nn.Module):
+    """Generic transformer-based binary classification model.
+
+    Args:
+        model_name (str): Name or path of the pretrained Hugging Face model.
+        use_mean_pooling (bool): If True, average token embeddings using the attention mask.
+                                 Otherwise, use pooler output or CLS token embedding.
+        dropout (float): Dropout rate applied before the final linear layer.
+                         Set to 0.0 to disable dropout. Must be ≥ 0.
+
+    pragma: skip_doc
+    """
+
+    MODEL_CLASS_PREFIXES: List[Tuple[str, str]] = [
+        ("bert-base-uncased", "BertModel"),
+        ("roberta-base", "RobertaModel"),
+        ("distilbert-base-uncased", "DistilBertModel"),
+        ("albert-base-v2", "AlbertModel"),
+        ("google/electra-base-discriminator", "ElectraModel"),
+    ]
+
+    def __init__(self,
+                 model_name: str,
+                 use_mean_pooling: bool = False,
+                 dropout: float = 0.1):
+        super().__init__()
+        if dropout < 0.0:
+            raise ValueError(f"Dropout must be non-negative, got {dropout}")
+
+        self.config = AutoConfig.from_pretrained(model_name)
+        model_class = self._resolve_model_class(model_name)
+        self.llm_model: PreTrainedModel = model_class.from_pretrained(model_name, config=self.config)
+
+        self.hidden_size = self.config.hidden_size
+        self.use_mean_pooling = use_mean_pooling
+        self.linear = nn.Linear(self.hidden_size, 1)
+
+        self.use_dropout = dropout > 0.0
+        if self.use_dropout:
+            self.dropout = nn.Dropout(dropout)
+
+        # Initialize linear weight
+        nn.init.xavier_uniform_(self.linear.weight)
+        if self.linear.bias is not None:
+            nn.init.zeros_(self.linear.bias)
+
+    def _resolve_model_class(self, model_name: str) -> Type[PreTrainedModel]:
+        for prefix, class_name in self.MODEL_CLASS_PREFIXES:
+            if model_name.startswith(prefix):
+                module = __import__("transformers", fromlist=[class_name])
+                return getattr(module, class_name)
+        raise ValueError(f"Model '{model_name}' is not supported.")
+
+    def forward(self, inputs:Tuple[torch.Tensor]) -> torch.Tensor:
+        """
+        Runs a forward pass through the transformer model followed by optional dropout and a classification head.
+
+        Args:
+            inputs (Tuple[torch.Tensor]): Tuple of below tensors:
+                input_ids (torch.Tensor): Token IDs of shape (batch_size, seq_len).
+                attention_mask (torch.Tensor): Attention mask of shape (batch_size, seq_len).
+
+        Returns:
+            torch.Tensor: Raw classification logits of shape (batch_size, 1).
+        """
+        input_ids, attention_mask = inputs
+        outputs = self.llm_model(input_ids=input_ids, attention_mask=attention_mask)
+
+        if self.use_mean_pooling:
+            last_hidden = outputs.last_hidden_state
+            # Expand the mask across the hidden dimension
+            # e.g. if hidden size is 768, the (seq_len,) mask becomes (seq_len, 768) per sample.
+            input_mask_expanded = attention_mask.unsqueeze(-1).expand(last_hidden.size())
+            sum_embeddings = torch.sum(last_hidden * input_mask_expanded, dim=1)
+            mask_adjusted_count = torch.clamp(input_mask_expanded.sum(1), min=1e-9)
+            x = sum_embeddings / mask_adjusted_count
+        elif hasattr(outputs, "pooler_output") and outputs.pooler_output is not None:
+            x = outputs.pooler_output
+        else:  # Grab the first feature in the sequence
+            x = outputs.last_hidden_state[:, 0, :]
+
+        assert x.dim() == 2
+
+        if self.use_dropout:
+            x = self.dropout(x)
+        x = self.linear(x)
+        return x

modules/pytorch_wrapper/ptn_lstm.py

@@ -36,7 +36,7 @@ class PtnLSTM:
                 "num_layers": ("INT", {"default": 1, "min": 1, "max": 1e3}),
                 "bias": ("BOOLEAN", {"default": True}),
                 "batch_first": ("BOOLEAN", {"default": True}),
-                "dropout": ("FLOAT", {"default": 0, "min": 0, "max": 1}),
+                "dropout": ("FLOAT", {"default": 0, "min": 0, "max": 1, "step": 0.001}),
                 "bidirectional": ("BOOLEAN", {"default": False}),
                 "proj_size": ("INT", {"default": 0, "min": 0, "max": 1e6}),
             }

modules/pytorch_wrapper/ptn_lstm_linear.py

@@ -108,7 +108,7 @@ class PtnLSTMLinear:
                 "num_layers": ("INT", {"default": 1, "min": 1, "max": 1e3}),
                 "bias": ("BOOLEAN", {"default": True}),
                 "batch_first": ("BOOLEAN", {"default": True}),
-                "dropout": ("FLOAT", {"default": 0, "min": 0, "max": 1}),
+                "dropout": ("FLOAT", {"default": 0, "min": 0, "max": 1, "step": 0.001}),
                 "bidirectional": ("BOOLEAN", {"default": False}),
                 "proj_size": ("INT", {"default": 0, "min": 0, "max": 1e6}),
                 "linear_output_size": ("INT", {"default": 1, "min": 1, "max": 1e6}),

modules/pytorch_wrapper/ptn_rnn.py

@@ -37,7 +37,7 @@ class PtnRNN:
                 "nonlinearity": (("tanh", "relu"),),
                 "bias": ("BOOLEAN", {"default": True}),
                 "batch_first": ("BOOLEAN", {"default": True}),
-                "dropout": ("FLOAT", {"default": 0, "min": 0, "max": 1}),
+                "dropout": ("FLOAT", {"default": 0, "min": 0, "max": 1, "step": 0.001}),
                 "bidirectional": ("BOOLEAN", {"default": False}),
             }
         }

modules/pytorch_wrapper/ptn_rnn_linear.py

@@ -123,7 +123,7 @@ class PtnRNNLinear:
                 "nonlinearity": (("tanh", "relu"),),
                 "bias": ("BOOLEAN", {"default": True}),
                 "batch_first": ("BOOLEAN", {"default": True}),
-                "dropout": ("FLOAT", {"default": 0, "min": 0, "max": 1}),
+                "dropout": ("FLOAT", {"default": 0, "min": 0, "max": 1, "step": 0.001}),
                 "bidirectional": ("BOOLEAN", {"default": False}),
                 "linear_output_size": ("INT", {"default": 1, "min": 1, "max": 1e6}),
                 "linear_bias": ("BOOLEAN", {"default": True}),

modules/pytorch_wrapper/ptv_hf_dataset_with_token_encode.py

@@ -1,6 +1,12 @@
+import os
+import sys
 from typing import Any, Dict, Tuple, Callable
 import torch
 
+PROJECT_ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", ".."))
+MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules")
+sys.path.append(MODULE_ROOT)
+from hugging_face_wrapper.hf_dataset_with_token_encode import HfDatasetWithTokenEncode
 
 class PtvHfDatasetWithTokenEncode:
     """
@@ -13,6 +19,7 @@ class PtvHfDatasetWithTokenEncode:
             label_field_name (str): Field name for labels.
             encode (PTCALLABLE): The reference to a token encoder function.
             remove_html_tags (bool): Remove html tags in text if True.
+            encode_return_dict (bool): Encode function returns a Dict instead of a Tuple.
 
     category: PyTorch wrapper - Training
     """
@@ -33,6 +40,7 @@ class PtvHfDatasetWithTokenEncode:
                 "label_field_name": ("STRING", {"default": "label","multiline": False}),
                 "encode": ("PTCALLABLE", {}),
                 "remove_html_tags": ("BOOLEAN", {"default": False}),
+                "encode_return_dict": ("BOOLEAN", {"default": False}),
             }
         }
 
@@ -50,7 +58,8 @@ class PtvHfDatasetWithTokenEncode:
           sample_field_name: str,
           label_field_name: str,
           encode: Callable,
-          remove_html_tags:bool) -> Tuple:
+          remove_html_tags:bool,
+          encode_return_dict:bool) -> Tuple:
         """
         Loads a dataset from Hugging Face with specified parameters.  
         
@@ -61,14 +70,11 @@ class PtvHfDatasetWithTokenEncode:
             label_field_name (str): Field name for labels.
             encode (Callable): The reference to a token encoder function.
             remove_html_tags (bool): Remove html tags in text if True.
+            encode_return_dict (bool): Encode function returns a Dict instead of a Tuple.
 
         Returns:  
             Tuple: A tuple containing the dataset instance.  
         """
-        # Keep import within this method so that relevant packages are imported
-        # only when they are actually needed.
-        from .hf_dataset_with_token_encode import HfDatasetWithTokenEncode
-
         with torch.inference_mode(False):
             dc = HfDatasetWithTokenEncode(
                 dataset_name=name,
@@ -76,5 +82,6 @@ class PtvHfDatasetWithTokenEncode:
                 sample_field_name=sample_field_name,
                 label_field_name=label_field_name,
                 encode=encode,
-                remove_html_tags=remove_html_tags)
+                remove_html_tags=remove_html_tags,
+                encode_return_dict=encode_return_dict)
             return (dc,)

modules/pytorch_wrapper/ptv_hf_glove_dataset.py

@@ -1,6 +1,12 @@
+import os
+import sys
 from typing import Any, Dict, Tuple
 import torch
 
+PROJECT_ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", ".."))
+MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules")
+sys.path.append(MODULE_ROOT)
+from hugging_face_wrapper.hf_glove_dataset import HfGloveDataset
 
 class PtvHfGloveDataset:
     """
@@ -65,9 +71,6 @@ class PtvHfGloveDataset:
         Returns:  
             Tuple: A tuple containing the dataset instance.  
         """
-        # Keep import within this method so that relevant packages are imported
-        # only when they are actually needed.
-        from .hf_glove_dataset import HfGloveDataset
 
         with torch.inference_mode(False):
             dc = HfGloveDataset(

modules/pytorch_wrapper/utils.py

@@ -4,7 +4,7 @@ from typing import Optional, Tuple, Union, List
 import random
 import numpy as np
 import torch
-
+import torch.nn as nn
 
 DTYPE_MAPPING = {
     "float32": torch.float32,
@@ -237,3 +237,16 @@ def pad_truncate_sequence(
         return torch.stack(token_list), torch.stack(mask_list)
     else:
         return token_list, mask_list
+
+
+def set_requires_grad(module: nn.Module, requires_grad: bool) -> None:
+    """
+    Enables or disables gradient computation for all parameters in a given module.
+
+    Args:
+        module (nn.Module): The PyTorch module whose parameters should be modified.
+        requires_grad (bool): If False, freezes the module (no gradients computed).
+                              If True, unfreezes the module (gradients computed).
+    """
+    for param in module.parameters():
+        param.requires_grad = requires_grad

pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "pt-wrapper"
 description = "PyTorch extension for ComfyUI featuring extensive PyTorch wrapper nodes for seamless tensor operations and PyTorch model training."
-version = "3.4.0"
+version = "3.5.0"
 license = {file = "LICENSE"}
 
 [project.urls]

requirements.txt

@@ -4,5 +4,6 @@ seaborn>=0.13.2
 scipy>=1.14.1
 scikit-learn>=1.6.1
 datasets>=3.4.1
+transformers>=4.47.1
 gensim==4.3.2
 sentencepiece>=0.2.0

tests/hugging_face_wrapper/hf_dataset_with_tokenize_test.py

@@ -0,0 +1,89 @@
+import os
+import sys
+import unittest
+import torch
+
+PROJECT_ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", ".."))
+MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules")
+sys.path.append(MODULE_ROOT)
+
+from hugging_face_wrapper.hf_dataset_with_token_encode import HfDatasetWithTokenEncode
+from sentencepiece_wrapper.sp_load_model import SpLoadModel
+from sentencepiece_wrapper.sp_encode import SpEncode
+
+
+class TestHfDatasetWithTokenEncode(unittest.TestCase):
+    """
+    Tests Hugging Face Dataset wrapper with custom token_encode function.
+    """
+    
+    def setUp(self):
+        """Set up test instance."""
+        self.load_node = SpLoadModel()
+        self.sp = self.load_node.f("spiece.model")[0]
+        self.encode_node = SpEncode()
+        self.max_length=15
+        encode = self.encode_node.f(
+            spmodel=self.sp,
+            padding=False,
+            padding_method="longest",
+            padding_value=0,
+            truncation=True,
+            max_length=self.max_length
+        )[0]
+
+        self.ds = HfDatasetWithTokenEncode(
+            "imdb", # dataset_name
+            "train", #split
+            encode=encode,
+        )
+        
+        self.ds_remove_html = HfDatasetWithTokenEncode(
+            "imdb", # dataset_name
+            "train", #split
+            encode=encode,
+            remove_html_tags=True
+        )
+
+    def test_loading(self):
+        it = iter(self.ds)
+        (sample, mask), label = next(it)
+
+        # Compare text
+        actual = sample.size()
+        expected = torch.Size((self.max_length, ))
+        self.assertEqual(expected, actual, f"expected {expected} and actual {actual} do not match.")
+
+        # Compare mask
+        actual = mask.size()
+        expected = torch.Size((self.max_length, ))
+        self.assertEqual(expected, actual, f"expected {expected} and actual {actual} do not match.")
+
+        # Compare label
+        actual = label
+        expected = torch.tensor(0, dtype=torch.int64)
+        self.assertEqual(expected, actual, f"expected {expected} and actual {actual} do not match.")
+
+
+    def test_loading_remove_html(self):
+        it = iter(self.ds_remove_html)
+        (sample, mask), label = next(it)
+
+        # Compare text
+        actual = sample.size()
+        expected = torch.Size((self.max_length, ))
+        self.assertEqual(expected, actual, f"expected {expected} and actual {actual} do not match.")
+
+        # Compare mask
+        actual = mask.size()
+        expected = torch.Size((self.max_length, ))
+        self.assertEqual(expected, actual, f"expected {expected} and actual {actual} do not match.")
+
+        # Compare label
+        actual = label
+        expected = torch.tensor(0, dtype=torch.int64)
+        self.assertEqual(expected, actual, f"expected {expected} and actual {actual} do not match.")
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/hugging_face_wrapper/hf_dataset_with_tokenize_with_hf_tokenize_encode_test.py

@@ -0,0 +1,86 @@
+import os
+import sys
+import unittest
+import torch
+
+PROJECT_ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", ".."))
+MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules")
+sys.path.append(MODULE_ROOT)
+
+from hugging_face_wrapper.hf_dataset_with_token_encode import HfDatasetWithTokenEncode
+from hugging_face_wrapper.hf_tokenizer_encode import HfTokenizerEncode
+
+
+class TestHfDatasetWithTokenEncode(unittest.TestCase):
+    """
+    Tests Hugging Face Dataset wrapper with custom token_encode function.
+    """
+    
+    def setUp(self):
+        """Set up test instance."""
+        self.max_length=15
+        encode = HfTokenizerEncode().f(
+            model_name="distilbert-base-uncased",
+            padding=True,
+            padding_method="longest",
+            truncation=True,
+            max_length=self.max_length
+        )[0]
+
+        self.ds = HfDatasetWithTokenEncode(
+            "imdb", # dataset_name
+            "train", #split
+            encode=encode,
+            encode_return_dict=True
+        )
+        
+        self.ds_remove_html = HfDatasetWithTokenEncode(
+            "imdb", # dataset_name
+            "train", #split
+            encode=encode,
+            remove_html_tags=True,
+            encode_return_dict=True
+        )
+
+    def test_loading(self):
+        it = iter(self.ds)
+        (sample, mask), label = next(it)
+
+        # Compare text
+        actual = sample.size()
+        expected = torch.Size((self.max_length,))
+        self.assertEqual(expected, actual, f"expected {expected} and actual {actual} do not match.")
+
+        # Compare mask
+        actual = mask.size()
+        expected = torch.Size((self.max_length,))
+        self.assertEqual(expected, actual, f"expected {expected} and actual {actual} do not match.")
+
+        # Compare label
+        actual = label
+        expected = torch.tensor(0, dtype=torch.int64)
+        self.assertEqual(expected, actual, f"expected {expected} and actual {actual} do not match.")
+
+
+    def test_loading_remove_html(self):
+        it = iter(self.ds_remove_html)
+        (sample, mask), label = next(it)
+
+        # Compare text
+        actual = sample.size()
+        expected = torch.Size((self.max_length, ))
+        self.assertEqual(expected, actual, f"expected {expected} and actual {actual} do not match.")
+
+        # Compare mask
+        actual = mask.size()
+        expected = torch.Size((self.max_length, ))
+        self.assertEqual(expected, actual, f"expected {expected} and actual {actual} do not match.")
+
+        # Compare label
+        actual = label
+        expected = torch.tensor(0, dtype=torch.int64)
+        self.assertEqual(expected, actual, f"expected {expected} and actual {actual} do not match.")
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/hugging_face_wrapper/hf_glove_dataset_test.py

@@ -0,0 +1,64 @@
+import os
+import sys
+import unittest
+import torch
+
+PROJECT_ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", ".."))
+MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules")
+sys.path.append(MODULE_ROOT)
+
+from pytorch_wrapper.hf_glove_dataset import HfGloveDataset
+
+
+class TestHfGloveDataset(unittest.TestCase):
+    """
+    Tests Hugging Face Dataset wrapper with word embedding support.
+    """
+    
+    def setUp(self):
+        """Set up test instance."""
+        self.embedding_dim = 100
+        self.max_seq_len = 256
+        self.ds = HfGloveDataset(
+            "imdb", # dataset_name
+            "train", #split
+            self.embedding_dim,
+            self.max_seq_len
+        )
+
+    def test_vocabulary_size(self):
+        self.assertEqual(400000, len(self.ds.glove))
+
+    def test_tokenize(self):
+        """Test tokenization"""
+        s = "Maddie is a cute puppy."
+        actual = self.ds.tokenize(s)
+        expected = ['maddie', 'is', 'a', 'cute', 'puppy']
+        self.assertEqual(expected, actual, f"expected {expected} and actual {actual} do not match.")
+
+    def test_vectorize(self):
+        """Test embedding generation"""
+        s = "Maddie is a cute puppy."
+        tokens = self.ds.tokenize(s)
+        embeddings = self.ds.vectorize(tokens)
+        actual = embeddings.size()
+        expected = torch.Size((self.max_seq_len, self.embedding_dim))
+        self.assertEqual(expected, actual, f"expected {expected} and actual {actual} do not match.")
+
+    def test_loading(self):
+        it = iter(self.ds)
+        sample = next(it)
+        
+        # Compare text
+        actual = sample[0].size()
+        expected = torch.Size((self.max_seq_len, self.embedding_dim))
+        self.assertEqual(expected, actual, f"expected {expected} and actual {actual} do not match.")
+
+        # Compare label
+        actual = sample[1]
+        expected = torch.tensor(0, dtype=torch.int64)
+        self.assertEqual(expected, actual, f"expected {expected} and actual {actual} do not match.")
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/hugging_face_wrapper/hf_tokenizer_encode_test.py

@@ -0,0 +1,97 @@
+import os
+import sys
+import unittest
+import torch
+from itertools import product
+
+PROJECT_ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", ".."))
+MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules")
+sys.path.append(MODULE_ROOT)
+
+from hugging_face_wrapper.hf_tokenizer_encode import HfTokenizerEncode
+
+
+class TestHfTokenizerEncode(unittest.TestCase):
+    def setUp(self):
+        self.node = HfTokenizerEncode()
+        self.test_sentence_1d = "Hello, world"
+        self.test_sentence_2d = [
+            "a b c d e f g h i j k l m n o p q r s t u v w x y z",
+            "a b c d e f g h i j k l"]
+        self.models = {
+            "bert-base-uncased": "BertTokenizer",
+            "roberta-base": "RobertaTokenizer",
+            "distilbert-base-uncased": "DistilBertTokenizer",
+            "albert-base-v2": "AlbertTokenizer",
+            "google/electra-small-discriminator": "ElectraTokenizer",
+        }
+
+    def test_tokenizer_variants(self):
+        for model_name, tokenizer_name in self.models.items():
+            for is_1d, padding, truncation, pad_method in product(
+                [True, False], [True, False], [True, False], ["longest", "max_length"]
+            ):
+                
+                if not is_1d and not padding:
+                    continue  # Cannot batch ragged output into tensors
+
+                with self.subTest(
+                    tokenizer_name=tokenizer_name,
+                    is_1d=is_1d,
+                    padding=padding,
+                    truncation=truncation,
+                    pad_method=pad_method
+                ):
+                    sentence = self.test_sentence_1d if is_1d else self.test_sentence_2d
+                    
+                    if truncation:
+                        max_length = 10
+                    else:
+                        max_length = 256
+
+                    encoder = self.node.f(
+                        model_name=model_name,
+                        padding=padding,
+                        padding_method=pad_method,
+                        truncation=truncation,
+                        max_length=max_length
+                    )[0]
+                    output = encoder(sentence=sentence)
+
+                    input_ids = output["input_ids"]
+                    mask = output["attention_mask"]
+
+                    # shape agreement
+                    self.assertEqual(input_ids.shape, mask.shape)
+
+                    # correct batch size
+                    self.assertEqual(input_ids.shape[0], 1 if is_1d else len(sentence))
+
+                    # correct max length behavior
+                    if not is_1d:
+                        if truncation:
+                            self.assertEqual(input_ids.shape[1], 10)
+                        elif padding and pad_method == "longest":
+                            self.assertEqual(input_ids[0].size(0), input_ids[1].size(0))
+                        elif padding and pad_method == "max_length":
+                            self.assertGreaterEqual(input_ids.shape[1], 256)
+
+
+    def test_padding_false_1d(self):
+        for model_name in self.models:
+            with self.subTest(model=model_name):
+                encoder = self.node.f(
+                    model_name=model_name,
+                    padding=False,
+                    padding_method="longest",
+                    truncation=False,
+                    max_length=256
+                )[0]
+                output = encoder(sentence=self.test_sentence_1d)
+                input_ids = output["input_ids"]
+                self.assertEqual(input_ids.shape[0], 1)
+                self.assertGreater(torch.count_nonzero(input_ids), 0)
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/pytorch_wrapper/pt_train_embedding_rnn_linear_text_classification_model_test.py

@@ -14,7 +14,7 @@ MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules")
 sys.path.append(MODULE_ROOT)
 
 from pytorch_wrapper.ptn_embedding_rnn_linear import PtnEmbeddingRNNLinear
-from pytorch_wrapper.hf_dataset_with_token_encode import HfDatasetWithTokenEncode
+from hugging_face_wrapper.hf_dataset_with_token_encode import HfDatasetWithTokenEncode
 from sentencepiece_wrapper.sp_load_model import SpLoadModel
 from sentencepiece_wrapper.sp_encode import SpEncode
 from pytorch_wrapper.pto_adam import PtoAdam

tests/pytorch_wrapper/pt_train_embedding_transformer_linear_text_classification_model_test.py

@@ -14,7 +14,7 @@ MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules")
 sys.path.append(MODULE_ROOT)
 
 from pytorch_wrapper.ptn_embedding_transformer_linear import PtnEmbeddingTransformerLinear
-from pytorch_wrapper.hf_dataset_with_token_encode import HfDatasetWithTokenEncode
+from hugging_face_wrapper.hf_dataset_with_token_encode import HfDatasetWithTokenEncode
 from sentencepiece_wrapper.sp_load_model import SpLoadModel
 from sentencepiece_wrapper.sp_encode import SpEncode
 from pytorch_wrapper.pto_adamw import PtoAdamW
@@ -105,7 +105,7 @@ class TestPtTrainEmbeddingTransformerLinearTextClassificationModel(unittest.Test
                                       )[0]
 
         self.minimum_lr = 1e-9
-        self.epochs = 10
+        self.epochs = 5
         self.scheduler = PtoLrSchedulerCosineAnnealing().f(self.optimizer,
                                                            self.epochs,
                                                            self.minimum_lr)[0]

tests/pytorch_wrapper/pt_train_fine_tune_classification_transformer_model_text_test.py

@@ -0,0 +1,139 @@
+"""
+Dataset URL:https://huggingface.co/datasets/stanfordnlp/imdb
+
+This unit test covers a test case where pretrained module is not frozen.
+For a case where pretrained module is frozen, see:
+tests/pytorch_wrapper/pt_train_fine_tune_classification_transformer_model_text_test.py
+
+For DistilBERT, this script should produce validation accuracy around 0.9242 for the first epoch and around 0.9216 for the secod epoch.
+"""
+import os
+import sys
+import unittest
+import torch
+import torch.nn as nn
+import random
+import numpy as np
+
+PROJECT_ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", ".."))
+MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules")
+sys.path.append(MODULE_ROOT)
+
+from pytorch_wrapper.ptn_hf_fine_tuned_classification_model import PtnHfFineTunedClassificationModel
+from hugging_face_wrapper.hf_dataset_with_token_encode import HfDatasetWithTokenEncode
+from hugging_face_wrapper.hf_tokenizer_encode import HfTokenizerEncode
+from pytorch_wrapper.pto_adamw import PtoAdamW
+from pytorch_wrapper.pt_data_loader import PtDataLoader
+from pytorch_wrapper.pt_train_classification_transformer_model import PtTrainClassificationTransformerModel
+from pytorch_wrapper.pt_save_model import PtSaveModel
+from pytorch_wrapper.pto_lr_scheduler_cosine_annealing import PtoLrSchedulerCosineAnnealing
+from pytorch_wrapper.ptn_bce_with_logits_loss import PtnBCEWithLogitsLoss
+from pytorch_wrapper.utils import set_seed
+
+
+class TestPtnHfFineTunedTextClassificationModel(unittest.TestCase):
+
+    @unittest.skipIf(os.getenv("RUN_SKIPPED_TESTS") != "1", "Skipping unless explicitly enabled")
+    def setUp(self):
+        set_seed(42)
+
+        self.loss_fun = PtnBCEWithLogitsLoss().f(reduction="mean")[0]
+        self.batch_size = 32
+        self.epochs = 2
+        self.minimum_lr = 1e-9
+        self.models = {
+            "bert-base-uncased": "BertTokenizer",
+            "distilbert-base-uncased": "DistilBertTokenizer",
+            "albert-base-v2": "AlbertTokenizer",
+            "roberta-base": "RobertaTokenizer",
+            # "google/electra-small-discriminator": "ElectraTokenizer",
+        }
+        self.model_path = "hf_fine_tuned_text_classification_model_test.pt"
+        self.linear_head = True
+
+    @unittest.skipIf(os.getenv("RUN_SKIPPED_TESTS") != "1", "Skipping unless explicitly enabled")
+    def test_all_models(self):
+        for model_name in self.models:
+            print(f"\n=== Testing model: {model_name} ===")
+
+            encode = HfTokenizerEncode().f(
+                model_name=model_name,
+                padding=True,
+                padding_method="max_length",
+                truncation=True,
+                max_length=0  # Use model's max length
+            )[0]
+
+            train_dataset = HfDatasetWithTokenEncode(
+                "imdb", "train", encode=encode,
+                remove_html_tags=True,
+                encode_return_dict=True
+            )
+
+            test_dataset = HfDatasetWithTokenEncode(
+                "imdb", "test", encode=encode,
+                remove_html_tags=True,
+                encode_return_dict=True
+            )
+
+            train_loader = PtDataLoader().f(
+                dataset=train_dataset,
+                batch_size=self.batch_size,
+                shuffle=True,
+                parameters='{"num_workers":1}'
+            )[0]
+
+            test_loader = PtDataLoader().f(
+                dataset=test_dataset,
+                batch_size=self.batch_size,
+                shuffle=False,
+                parameters='{"num_workers":1}'
+            )[0]
+
+            model = PtnHfFineTunedClassificationModel().f(
+                model_name,
+                use_mean_pooling=True,
+                dropout=0.1
+            )[0]
+
+            optimizer = PtoAdamW().f(
+                model, 1e-5, 0.9, 0.999,
+                0.01,  # weight decay
+                False  # amsgrad
+            )[0]
+
+            scheduler = PtoLrSchedulerCosineAnnealing().f(
+                optimizer,
+                self.epochs,
+                self.minimum_lr
+            )[0]
+
+            trainer = PtTrainClassificationTransformerModel()
+
+            trained_model = trainer.f(
+                model,
+                train_loader,
+                optimizer,
+                self.loss_fun,
+                self.epochs,
+                use_gpu=True,
+                early_stopping=False,
+                early_stopping_rounds=1,
+                scheduler=scheduler,
+                output_best_val_model=False,
+                classification_metrics=True,
+                val_loader=test_loader
+            )[0]
+
+            PtSaveModel().f(trained_model, self.model_path)
+
+            # Clean up to free memory and VRAM
+            del model, trained_model, optimizer, scheduler, trainer
+            del train_loader, test_loader, train_dataset, test_dataset, encode
+            torch.cuda.empty_cache()
+            import gc
+            gc.collect()
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/pytorch_wrapper/pt_train_gru_linear_text_classification_model_test.py

@@ -84,13 +84,13 @@ class TestPtTrainGRULinearTextClassificationModel(unittest.TestCase):
         )[0]
 
         self.optimizer = PtoAdamW().f(self.model,
-                                      0.001,
+                                      0.0003,
                                       0.9, 0.999,
                                       weight_decay=0.01,
                                       amsgrad=False)[0]
         
         self.scheduler = PtoLrSchedulerReduceOnPlateau().f(self.optimizer,
-                                                grace_period=10,
+                                                grace_period=5,
                                                 gamma=0.5
                                                )[0]
 
@@ -98,7 +98,7 @@ class TestPtTrainGRULinearTextClassificationModel(unittest.TestCase):
         self.save_model = PtSaveModel()
         self.load_model = PtLoadModel()
         self.linear_head = True
-        self.epochs = 50
+        self.epochs = 20
 
     @unittest.skipIf(os.getenv("RUN_SKIPPED_TESTS") != "1", "Skipping unless explicitly enabled")
     def test_1(self):

tests/pytorch_wrapper/pt_train_hf_fine_tuned_text_classification_model_test.py

@@ -0,0 +1,131 @@
+"""
+Dataset URL:https://huggingface.co/datasets/stanfordnlp/imdb
+
+This unit test covers a rare test case where pretrained module is not frozen.
+For a case where pretrained module is frozen, see:
+tests/pytorch_wrapper/pt_train_fine_tune_classification_transformer_model_text_test.py
+
+This script should produce validation accuracy around 0.9242 for the first epoch and around 0.9216 for the secod epoch.
+"""
+import os
+import sys
+import unittest
+import torch
+import torch.nn as nn
+import random
+import numpy as np
+
+PROJECT_ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", ".."))
+MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules")
+sys.path.append(MODULE_ROOT)
+
+from pytorch_wrapper.ptn_hf_fine_tuned_classification_model import PtnHfFineTunedClassificationModel
+from hugging_face_wrapper.hf_dataset_with_token_encode import HfDatasetWithTokenEncode
+from hugging_face_wrapper.hf_tokenizer_encode import HfTokenizerEncode
+from pytorch_wrapper.pto_adamw import PtoAdamW
+from pytorch_wrapper.pt_data_loader import PtDataLoader
+from pytorch_wrapper.pt_train_classification_transformer_model import PtTrainClassificationTransformerModel
+from pytorch_wrapper.pt_save_model import PtSaveModel
+from pytorch_wrapper.pt_load_model import PtLoadModel
+from pytorch_wrapper.pto_lr_scheduler_cosine_annealing import PtoLrSchedulerCosineAnnealing
+from pytorch_wrapper.ptn_bce_with_logits_loss import PtnBCEWithLogitsLoss
+from pytorch_wrapper.utils import set_seed
+
+
+class TestPtnHfFineTunedTextClassificationModel(unittest.TestCase):
+
+    @unittest.skipIf(os.getenv("RUN_SKIPPED_TESTS") != "1", "Skipping unless explicitly enabled")
+    def setUp(self):
+        set_seed(42)
+
+        self.loss_fun =PtnBCEWithLogitsLoss().f(reduction="mean")[0]
+
+        self.max_length=512
+        self.model_name = "distilbert-base-uncased"
+        encode = HfTokenizerEncode().f(
+            model_name=self.model_name,
+            padding=True,
+            padding_method="max_length",
+            truncation=True,
+            max_length=0  # Use model's max length
+        )[0]
+
+        train_dataset = HfDatasetWithTokenEncode(
+            "imdb", # dataset_name
+            "train", #split
+            encode=encode,
+            remove_html_tags=True,
+            encode_return_dict=True
+        )
+
+        test_dataset = HfDatasetWithTokenEncode(
+            "imdb", # dataset_name
+            "test", #split
+            encode=encode,
+            remove_html_tags=True,
+            encode_return_dict=True
+        )
+
+        train_data_loader_node = PtDataLoader()
+        test_data_loader_node = PtDataLoader()
+
+        self.train_loader = train_data_loader_node.f(
+            dataset=train_dataset,
+            batch_size=32,
+            shuffle=True,
+            parameters='{"num_workers":1}'
+        )[0]
+
+        self.test_loader = test_data_loader_node.f(
+            dataset=test_dataset,
+            batch_size=32,
+            shuffle=False,
+            parameters='{"num_workers":1}'
+        )[0]
+
+        self.model_path = "hf_fine_tuned_text_classification_model_test.pt"
+
+        self.model = PtnHfFineTunedClassificationModel().f(
+            self.model_name,
+            use_mean_pooling=True,
+            dropout=0.1
+        )[0]
+
+        self.optimizer = PtoAdamW().f(self.model, 1e-5, 0.9, 0.999,
+                                      0.01,  # weight decay
+                                      False  # amsgrad
+                                      )[0]
+
+        self.minimum_lr = 1e-9
+        self.epochs = 2
+        self.scheduler = PtoLrSchedulerCosineAnnealing().f(self.optimizer,
+                                                           self.epochs,
+                                                           self.minimum_lr)[0]
+
+        self.trainer = PtTrainClassificationTransformerModel()
+        self.save_model = PtSaveModel()
+        self.load_model = PtLoadModel()
+        self.linear_head = True
+
+
+    @unittest.skipIf(os.getenv("RUN_SKIPPED_TESTS") != "1", "Skipping unless explicitly enabled")
+    def test_1(self):
+
+        trained_model = self.trainer.f(
+                        self.model,
+                        self.train_loader,
+                        self.optimizer,
+                        self.loss_fun,
+                        self.epochs, # epochs
+                        use_gpu=True,
+                        early_stopping=False,
+                        early_stopping_rounds=1,
+                        scheduler=self.scheduler,
+                        output_best_val_model=False,
+                        classification_metrics=True,
+                        val_loader=self.test_loader)[0]
+        self.save_model.f(trained_model, self.model_path)
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/pytorch_wrapper/pt_train_lstm_linear_text_classification_model_test.py

@@ -85,13 +85,13 @@ class TestPtTrainLSTMLinearTextClassificationModel(unittest.TestCase):
         )[0]
 
         self.optimizer = PtoAdamW().f(self.model,
-                                      0.001,
+                                      0.0003,
                                       0.9, 0.999,
                                       weight_decay=0.01,
                                       amsgrad=False)[0]
         
         self.scheduler = PtoLrSchedulerReduceOnPlateau().f(self.optimizer,
-                                                grace_period=10,
+                                                grace_period=5,
                                                 gamma=0.5
                                                )[0]
 
@@ -99,7 +99,7 @@ class TestPtTrainLSTMLinearTextClassificationModel(unittest.TestCase):
         self.save_model = PtSaveModel()
         self.load_model = PtLoadModel()
         self.linear_head = True
-        self.epochs = 50
+        self.epochs = 15
 
     @unittest.skipIf(os.getenv("RUN_SKIPPED_TESTS") != "1", "Skipping unless explicitly enabled")
     def test_1(self):

tests/pytorch_wrapper/pt_train_rnn_linear_text_classification_model_test.py

@@ -95,7 +95,7 @@ class TestPtTrainRNNLinearTextClassificationModel(unittest.TestCase):
         self.save_model = PtSaveModel()
         self.load_model = PtLoadModel()
         self.linear_head = True
-        self.epochs = 50
+        self.epochs = 10
 
     @unittest.skipIf(os.getenv("RUN_SKIPPED_TESTS") != "1", "Skipping unless explicitly enabled")
     def test_1(self):

tests/pytorch_wrapper/ptn_hf_fine_tuned_classification_model_def_test.py

@@ -0,0 +1,52 @@
+import os
+import sys
+import unittest
+import torch
+
+PROJECT_ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", ".."))
+MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules")
+sys.path.append(MODULE_ROOT)
+from pytorch_wrapper.ptn_hf_fine_tuned_classification_model_def import HfFineTunedClassificationModel
+from hugging_face_wrapper.hf_tokenizer_encode import HfTokenizerEncode
+
+
+class TestHfFineTunedClassificationModel(unittest.TestCase):
+    def setUp(self):
+        """Setup shared test conditions."""
+        self.text = [
+            "I really liked the scenary.",
+            "I did not like the vision of the creator."
+        ]
+        self.model_names = [
+            "distilbert-base-uncased",
+            "bert-base-uncased",
+            "roberta-base",
+            "albert-base-v2",
+            "google/electra-base-discriminator"
+        ]
+
+    def test_all_models(self):
+        """Run output shape test across all supported models."""
+        for model_name in self.model_names:
+            with self.subTest(model_name=model_name):
+                encode = HfTokenizerEncode().f(
+                    model_name=model_name,
+                    padding=True,
+                    padding_method="max_length",
+                    truncation=True,
+                    max_length=0  # Use model's max length
+                )[0]
+
+                model = HfFineTunedClassificationModel(
+                    model_name,
+                    use_mean_pooling=True,
+                    dropout=0.1
+                )
+
+                token_dict = encode(self.text)
+                out = model((token_dict["input_ids"], token_dict["attention_mask"]))
+                self.assertEqual(out.size(), torch.Size([2, 1]))
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/pytorch_wrapper/ptn_hf_fine_tuned_classification_model_test.py

@@ -0,0 +1,52 @@
+import os
+import sys
+import unittest
+import torch
+
+PROJECT_ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", ".."))
+MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules")
+sys.path.append(MODULE_ROOT)
+from pytorch_wrapper.ptn_hf_fine_tuned_classification_model import PtnHfFineTunedClassificationModel
+from hugging_face_wrapper.hf_tokenizer_encode import HfTokenizerEncode
+
+
+class TestPtfHfFineTunedClassificationModel(unittest.TestCase):
+    def setUp(self):
+        """Setup shared test conditions."""
+        self.text = [
+            "I really liked the scenary.",
+            "I did not like the vision of the creator."
+        ]
+        self.model_names = [
+            "distilbert-base-uncased",
+            "bert-base-uncased",
+            "roberta-base",
+            "albert-base-v2",
+            "google/electra-base-discriminator"
+        ]
+
+    def test_all_models(self):
+        """Run output shape test across all supported models."""
+        for model_name in self.model_names:
+            with self.subTest(model_name=model_name):
+                encode = HfTokenizerEncode().f(
+                    model_name=model_name,
+                    padding=True,
+                    padding_method="max_length",
+                    truncation=True,
+                    max_length=0  # Use model's max length
+                )[0]
+
+                model = PtnHfFineTunedClassificationModel().f(
+                    model_name,
+                    use_mean_pooling=True,
+                    dropout=0.1
+                )[0]
+
+                token_dict = encode(self.text)
+                out = model((token_dict["input_ids"], token_dict["attention_mask"]))
+                self.assertEqual(out.size(), torch.Size([2, 1]))
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/pytorch_wrapper/ptv_hf_dataset_with_token_encode_test.py

@@ -37,7 +37,8 @@ class TestPtvHfDatasetWithTokenEncode(unittest.TestCase):
             "text",
             "label",
             encode,
-            False  # remove_html_tags
+            False,  # remove_html_tags
+            False   # encode_return_dict
         )[0]
 
 

tests/pytorch_wrapper/utils_set_requires_grad_test.py

@@ -0,0 +1,36 @@
+import os
+import sys
+import unittest
+import torch.nn as nn
+
+
+PROJECT_ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", ".."))
+MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules")
+sys.path.append(MODULE_ROOT)
+
+from pytorch_wrapper.utils import set_requires_grad
+
+
+class TestSetRequiresGrad(unittest.TestCase):
+    def setUp(self):
+        self.model = nn.Linear(2, 3, bias=True)
+
+    def test_1(self):
+        # Check that gradients are initially enabled
+        for param in self.model.parameters():
+            self.assertTrue(param.requires_grad)
+
+        # Turn off gradients
+        set_requires_grad(self.model, False)
+        for param in self.model.parameters():
+            self.assertFalse(param.requires_grad)
+
+        # Turn on gradients
+        set_requires_grad(self.model, True)
+        for param in self.model.parameters():
+            self.assertTrue(param.requires_grad)
+
+
+if __name__ == "__main__":
+    unittest.main()
+