From 312f25067262ebeca80c049d60d87ec9f4281558 Mon Sep 17 00:00:00 2001 From: HowToSD Date: Fri, 28 Mar 2025 19:16:55 -0700 Subject: [PATCH] Added Transformer support --- __init__.py | 6 + docs/reference/node_reference.md | 2 + ..._train_classification_transformer_model.md | 44 +++ .../ptn_embedding_transformer_linear.md | 46 ++++ .../ptv_hf_dataset_with_token_encode.md | 2 + .../embedding_transformer_classification.json | 1 + .../hf_dataset_with_token_encode.py | 27 +- ..._train_classification_transformer_model.py | 231 ++++++++++++++++ .../ptn_embedding_transformer_linear.py | 251 ++++++++++++++++++ .../pto_lr_scheduler_cosine_annealing.py | 2 +- .../ptv_hf_dataset_with_token_encode.py | 9 +- pyproject.toml | 2 +- .../hf_dataset_with_tokenize_test.py | 27 ++ ...r_linear_text_classification_model_test.py | 139 ++++++++++ .../ptn_embedding_transformer_linear_test.py | 43 +++ .../ptv_hf_dataset_with_token_encode_test.py | 3 +- 16 files changed, 826 insertions(+), 9 deletions(-) create mode 100644 docs/reference/pt_train_classification_transformer_model.md create mode 100644 docs/reference/ptn_embedding_transformer_linear.md create mode 100644 examples/workflows/embedding_transformer_classification.json create mode 100644 modules/pytorch_wrapper/pt_train_classification_transformer_model.py create mode 100644 modules/pytorch_wrapper/ptn_embedding_transformer_linear.py create mode 100644 tests/pytorch_wrapper/pt_train_embedding_transformer_linear_text_classification_model_test.py create mode 100644 tests/pytorch_wrapper/ptn_embedding_transformer_linear_test.py diff --git a/__init__.py b/__init__.py index 4d5ae3f..bab0007 100644 --- a/__init__.py +++ b/__init__.py @@ -110,6 +110,7 @@ from .modules.pytorch_wrapper.pt_to_uint8 import PtToUint8 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_model import PtTrainModel from .modules.pytorch_wrapper.pt_train_regression_model import PtTrainRegressionModel from .modules.pytorch_wrapper.pt_train_rnn_model import PtTrainRNNModel @@ -159,6 +160,7 @@ from .modules.pytorch_wrapper.ptn_conv_model import PtnConvModel from .modules.pytorch_wrapper.ptn_cross_entropy_loss import PtnCrossEntropyLoss from .modules.pytorch_wrapper.ptn_embedding import PtnEmbedding from .modules.pytorch_wrapper.ptn_embedding_rnn_linear import PtnEmbeddingRNNLinear +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_huber_loss import PtnHuberLoss @@ -317,6 +319,7 @@ NODE_CLASS_MAPPINGS: Dict[str, Type[T]] = { "PtTokenizer": PtTokenizer, "PtTrainClassificationModel": PtTrainClassificationModel, "PtTrainClassificationModelLr": PtTrainClassificationModelLr, + "PtTrainClassificationTransformerModel": PtTrainClassificationTransformerModel, "PtTrainModel": PtTrainModel, "PtTrainRNNModel": PtTrainRNNModel, "PtTrainRegressionModel": PtTrainRegressionModel, @@ -366,6 +369,7 @@ NODE_CLASS_MAPPINGS: Dict[str, Type[T]] = { "PtnCrossEntropyLoss": PtnCrossEntropyLoss, "PtnEmbedding": PtnEmbedding, "PtnEmbeddingRNNLinear": PtnEmbeddingRNNLinear, + "PtnEmbeddingTransformerLinear": PtnEmbeddingTransformerLinear, "PtnGRU": PtnGRU, "PtnGRULinear": PtnGRULinear, "PtnHuberLoss": PtnHuberLoss, @@ -524,6 +528,7 @@ NODE_DISPLAY_NAME_MAPPINGS: Dict[str, str] = { "PtTokenizer": "Pt Tokenizer", "PtTrainClassificationModel": "Pt Train Classification Model", "PtTrainClassificationModelLr": "Pt Train Classification Model Lr", + "PtTrainClassificationTransformerModel": "Pt Train Classification Transformer Model", "PtTrainModel": "Pt Train Model", "PtTrainRNNModel": "Pt Train RNN Model", "PtTrainRegressionModel": "Pt Train Regression Model", @@ -573,6 +578,7 @@ NODE_DISPLAY_NAME_MAPPINGS: Dict[str, str] = { "PtnCrossEntropyLoss": "Ptn Cross Entropy Loss", "PtnEmbedding": "Ptn Embedding", "PtnEmbeddingRNNLinear": "Ptn Embedding RNN Linear", + "PtnEmbeddingTransformerLinear": "Ptn Embedding Transformer Linear", "PtnGRU": "Ptn GRU", "PtnGRULinear": "Ptn GRU Linear", "PtnHuberLoss": "Ptn Huber Loss", diff --git a/docs/reference/node_reference.md b/docs/reference/node_reference.md index bcc50f0..9cc949a 100644 --- a/docs/reference/node_reference.md +++ b/docs/reference/node_reference.md @@ -157,6 +157,7 @@ | [Ptn Conv 2d](ptn_conv2d.md) | A convolutional model consisting of a single conv2d layer. | | [Ptn Conv Model](ptn_conv_model.md) | A convolutional model consisting of multiple convolutional layers. | | [Ptn Embedding](ptn_embedding.md) | Constructs an embedding layer. | +| [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 Instance Norm 2d](ptn_instance_norm2d.md) | A normalization model to normalize elements over spatial axes within each channel for each sample. | @@ -234,6 +235,7 @@ | [Pt Predict Regression Model](pt_predict_regression_model.md) | Performs inference on input data. | | [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 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. | diff --git a/docs/reference/pt_train_classification_transformer_model.md b/docs/reference/pt_train_classification_transformer_model.md new file mode 100644 index 0000000..a7259e5 --- /dev/null +++ b/docs/reference/pt_train_classification_transformer_model.md @@ -0,0 +1,44 @@ +# Pt Train Classification Transformer Model +Trains 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. + loss (PTLOSS): The loss class. + epochs (int): Number of training epochs. + 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). + scheduler (PTLRSCHEDULER): Learning rate scheduler. + +## Input +| Name | Data type | +|---|---| +| model | Ptmodel | +| train_loader | Ptdataloader | +| optimizer | Ptoptimizer | +| loss_function | Ptloss | +| 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 | + +
+Category: PyTorch wrapper - Training + +ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved. diff --git a/docs/reference/ptn_embedding_transformer_linear.md b/docs/reference/ptn_embedding_transformer_linear.md new file mode 100644 index 0000000..1d780c9 --- /dev/null +++ b/docs/reference/ptn_embedding_transformer_linear.md @@ -0,0 +1,46 @@ +# Ptn Embedding Transformer Linear +A Transformer encoder model with a linear head. + + Args: + num_layers (int): Number of encoder layers in Transformer. + vocabulary_size (int): The vocabulary size of the tokens. + hidden_size (int): Hidden dimension (d_model) for each token of the Transformer. + nhead (int): Number of attention heads. + dim_feedforward (int): Feedforward dimension. + dropout (float): Dropout rate for Transformer layers. + nonlinearity (str): Activation function ('relu' or 'gelu') used in Transformer layers. + layer_norm_eps (float): Epsilon value for layer normalization. + batch_first (bool): Input shape format: True for (B, Seq, Feature), False for (Seq, B, Feature). + norm_first (bool): Whether to apply normalization before attention/feedforward. + bias (bool): Use bias in the Transformer layers. + max_length (int): Maximum input sequence length (for positional embedding). + linear_output_size (int): Output size of the final linear layer. + linear_bias (bool): Use bias in the linear layer. + +## Input +| Name | Data type | +|---|---| +| num_layers | Int | +| vocabulary_size | Int | +| hidden_size | Int | +| nhead | Int | +| dim_feedforward | Int | +| dropout | Float | +| nonlinearity | | +| layer_norm_eps | Float | +| batch_first | Boolean | +| norm_first | Boolean | +| bias | Boolean | +| max_length | Int | +| linear_output_size | Int | +| linear_bias | Boolean | + +## Output +| Data type | +|---| +| Ptmodel | + +
+Category: PyTorch wrapper - Model + +ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved. diff --git a/docs/reference/ptv_hf_dataset_with_token_encode.md b/docs/reference/ptv_hf_dataset_with_token_encode.md index f1f27af..a89b877 100644 --- a/docs/reference/ptv_hf_dataset_with_token_encode.md +++ b/docs/reference/ptv_hf_dataset_with_token_encode.md @@ -7,6 +7,7 @@ A PyTorch Dataset class wrapper for a Hugging Face dataset that converts text to sample_field_name (str): Field name for text samples. 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. ## Input | Name | Data type | @@ -16,6 +17,7 @@ A PyTorch Dataset class wrapper for a Hugging Face dataset that converts text to | sample_field_name | String | | label_field_name | String | | encode | Ptcallable | +| remove_html_tags | Boolean | ## Output | Data type | diff --git a/examples/workflows/embedding_transformer_classification.json b/examples/workflows/embedding_transformer_classification.json new file mode 100644 index 0000000..b0161b2 --- /dev/null +++ b/examples/workflows/embedding_transformer_classification.json @@ -0,0 +1 @@ +{"last_node_id":29,"last_link_id":43,"nodes":[{"id":19,"type":"PandasShowText","pos":[3504.62841796875,1897.03662109375],"size":[302.7994079589844,112],"flags":{},"order":12,"mode":0,"inputs":[{"name":"text","type":"STRING","widget":{"name":"text"},"link":20}],"outputs":[{"name":"STRING","localized_name":"STRING","type":"STRING","shape":6,"links":null}],"properties":{"Node name for S&R":"PandasShowText"},"widgets_values":["","/container/transfer/models/embedding_transformer_classification.pt"]},{"id":25,"type":"SpLoadModel","pos":[1246.61181640625,1912.7467041015625],"size":[210,60.936126708984375],"flags":{},"order":0,"mode":0,"inputs":[],"outputs":[{"name":"spmodel","localized_name":"spmodel","type":"SPMODEL","links":[26],"slot_index":0}],"properties":{"Node name for S&R":"SpLoadModel"},"widgets_values":["spiece.model"]},{"id":21,"type":"PtvHfDatasetWithTokenEncode","pos":[1907.6329345703125,1746.281982421875],"size":[271.5191650390625,154],"flags":{},"order":5,"mode":0,"inputs":[{"name":"encode","localized_name":"encode","type":"PTCALLABLE","link":27}],"outputs":[{"name":"PTVDATASET","localized_name":"PTVDATASET","type":"PTVDATASET","links":[24],"slot_index":0}],"properties":{"Node name for S&R":"PtvHfDatasetWithTokenEncode"},"widgets_values":["imdb","train","text","label",true]},{"id":24,"type":"PtvHfDatasetWithTokenEncode","pos":[1911.3275146484375,2022.5567626953125],"size":[271.5191650390625,154],"flags":{},"order":6,"mode":0,"inputs":[{"name":"encode","localized_name":"encode","type":"PTCALLABLE","link":28}],"outputs":[{"name":"PTVDATASET","localized_name":"PTVDATASET","type":"PTVDATASET","links":[25],"slot_index":0}],"properties":{"Node name for S&R":"PtvHfDatasetWithTokenEncode"},"widgets_values":["imdb","test","text","label",true]},{"id":16,"type":"PtnBCEWithLogitsLoss","pos":[2668.492919921875,1800.2408447265625],"size":[315,58],"flags":{},"order":1,"mode":0,"inputs":[],"outputs":[{"name":"PTLOSS","localized_name":"PTLOSS","type":"PTLOSS","links":[36],"slot_index":0}],"properties":{"Node name for S&R":"PtnBCEWithLogitsLoss"},"widgets_values":["mean"]},{"id":7,"type":"PtDataLoader","pos":[2293.7060546875,1735.56494140625],"size":[271.68658447265625,173.3039093017578],"flags":{},"order":8,"mode":0,"inputs":[{"name":"dataset","localized_name":"dataset","type":"PTVDATASET","link":24}],"outputs":[{"name":"PTDATALOADER","localized_name":"PTDATALOADER","type":"PTDATALOADER","links":[34],"slot_index":0}],"properties":{"Node name for S&R":"PtDataLoader"},"widgets_values":[32,true,"{\"num_workers\":1}"]},{"id":12,"type":"PtDataLoader","pos":[2290.658447265625,2016.2130126953125],"size":[271.68658447265625,173.3039093017578],"flags":{},"order":9,"mode":0,"inputs":[{"name":"dataset","localized_name":"dataset","type":"PTVDATASET","link":25}],"outputs":[{"name":"PTDATALOADER","localized_name":"PTDATALOADER","type":"PTDATALOADER","links":[35],"slot_index":0}],"properties":{"Node name for S&R":"PtDataLoader"},"widgets_values":[32,false,"{\"num_workers\":1}"]},{"id":17,"type":"PtSaveModel","pos":[3495.453125,1735.56494140625],"size":[318.3275146484375,91.27499389648438],"flags":{},"order":11,"mode":0,"inputs":[{"name":"model","localized_name":"model","type":"PTMODEL","link":38}],"outputs":[{"name":"model full path","localized_name":"model full path","type":"STRING","links":[20],"slot_index":0}],"properties":{"Node name for S&R":"PtSaveModel"},"widgets_values":["/container/transfer/models/embedding_transformer_classification.pt"]},{"id":28,"type":"PtoLrSchedulerCosineAnnealing","pos":[3098.5263671875,1348.582275390625],"size":[311.11163330078125,113.72386932373047],"flags":{},"order":7,"mode":0,"inputs":[{"name":"optimizer","localized_name":"optimizer","type":"PTOPTIMIZER","link":43}],"outputs":[{"name":"PTLRSCHEDULER","localized_name":"PTLRSCHEDULER","type":"PTLRSCHEDULER","links":[40],"slot_index":0}],"properties":{"Node name for S&R":"PtoLrSchedulerCosineAnnealing"},"widgets_values":[10,1.0000000000000005e-9]},{"id":29,"type":"PtoAdamW","pos":[2673.749267578125,1107.9647216796875],"size":[285.66900634765625,158.52035522460938],"flags":{},"order":4,"mode":0,"inputs":[{"name":"model","localized_name":"model","type":"PTMODEL","link":41}],"outputs":[{"name":"PTOPTIMIZER","localized_name":"PTOPTIMIZER","type":"PTOPTIMIZER","links":[42,43],"slot_index":0}],"properties":{"Node name for S&R":"PtoAdamW"},"widgets_values":[0.0003000000000000002,0.9,0.999,0.01,false]},{"id":26,"type":"PtnEmbeddingTransformerLinear","pos":[2669.3544921875,1336.1724853515625],"size":[303.65289306640625,370],"flags":{},"order":2,"mode":0,"inputs":[],"outputs":[{"name":"PTMODEL","localized_name":"PTMODEL","type":"PTMODEL","links":[32,41],"slot_index":0}],"properties":{"Node name for S&R":"PtnEmbeddingTransformerLinear"},"widgets_values":[2,32000,256,8,512,0.5,"gelu",0.00001,true,false,true,512,1,true]},{"id":27,"type":"PtTrainClassificationTransformerModel","pos":[3090.223876953125,1736.2188720703125],"size":[344.3999938964844,292.4398193359375],"flags":{},"order":10,"mode":0,"inputs":[{"name":"model","localized_name":"model","type":"PTMODEL","link":32},{"name":"train_loader","localized_name":"train_loader","type":"PTDATALOADER","link":34},{"name":"optimizer","localized_name":"optimizer","type":"PTOPTIMIZER","link":42},{"name":"loss_function","localized_name":"loss_function","type":"PTLOSS","link":36},{"name":"scheduler","localized_name":"scheduler","type":"PTLRSCHEDULER","shape":7,"link":40},{"name":"val_loader","localized_name":"val_loader","type":"PTDATALOADER","shape":7,"link":35}],"outputs":[{"name":"Model","localized_name":"Model","type":"PTMODEL","links":[38],"slot_index":0},{"name":"train loss","localized_name":"train loss","type":"TENSOR","links":null},{"name":"val loss","localized_name":"val loss","type":"TENSOR","links":null}],"properties":{"Node name for S&R":"PtTrainClassificationTransformerModel"},"widgets_values":[10,true,false,10,false,true]},{"id":22,"type":"SpEncode","pos":[1502.9549560546875,1913.823974609375],"size":[315,154],"flags":{},"order":3,"mode":0,"inputs":[{"name":"spmodel","localized_name":"spmodel","type":"SPMODEL","link":26}],"outputs":[{"name":"PTCALLABLE","localized_name":"PTCALLABLE","type":"PTCALLABLE","links":[27,28],"slot_index":0}],"properties":{"Node name for S&R":"SpEncode"},"widgets_values":[true,"max_length",0,true,512]}],"links":[[20,17,0,19,0,"STRING"],[24,21,0,7,0,"PTVDATASET"],[25,24,0,12,0,"PTVDATASET"],[26,25,0,22,0,"SPMODEL"],[27,22,0,21,0,"PTCALLABLE"],[28,22,0,24,0,"PTCALLABLE"],[32,26,0,27,0,"PTMODEL"],[34,7,0,27,1,"PTDATALOADER"],[35,12,0,27,5,"PTDATALOADER"],[36,16,0,27,3,"PTLOSS"],[38,27,0,17,0,"PTMODEL"],[40,28,0,27,4,"PTLRSCHEDULER"],[41,26,0,29,0,"PTMODEL"],[42,29,0,27,2,"PTOPTIMIZER"],[43,29,0,28,0,"PTOPTIMIZER"]],"groups":[],"config":{},"extra":{"ds":{"scale":1.115326655450681,"offset":[-1077.6725411509767,-1390.0549035093902]},"node_versions":{"data-analysis":"68dd090367adffa569c3b3993d56e422f4516d2d","pt-wrapper":"3ba2eb708ba62ffb55b9855d79346746b961c250"},"ue_links":[]},"version":0.4} \ No newline at end of file diff --git a/modules/pytorch_wrapper/hf_dataset_with_token_encode.py b/modules/pytorch_wrapper/hf_dataset_with_token_encode.py index b67e799..0db4df3 100644 --- a/modules/pytorch_wrapper/hf_dataset_with_token_encode.py +++ b/modules/pytorch_wrapper/hf_dataset_with_token_encode.py @@ -1,4 +1,6 @@ -from typing import Tuple, Callable +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) @@ -15,7 +17,8 @@ class HfDatasetWithTokenEncode(Dataset): split: str, sample_field_name: str = "text", label_field_name: str = "label", - encode: Callable=None) -> None: + encode: Optional[Callable]=None, + remove_html_tags:Optional[bool]=False) -> None: """ Initializes the dataset by loading Hugging Face dataset. @@ -24,12 +27,14 @@ class HfDatasetWithTokenEncode(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 (Callable): Token encode function + encode (Optional[Callable]): Token encode function + remove_html_tags (Optional[bool]): Remove html tags in text if True. """ 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 def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor]: @@ -46,7 +51,21 @@ class HfDatasetWithTokenEncode(Dataset): """ # Get a text sample from database and convert to a word embedding sample = self.dataset[idx] - tokens, masks = self.encode(sample[self.sample_field_name]) + 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
,

, 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] + + tokens, masks = self.encode(text) # Get a corresponding label label = torch.tensor(sample[self.label_field_name], dtype=torch.long) diff --git a/modules/pytorch_wrapper/pt_train_classification_transformer_model.py b/modules/pytorch_wrapper/pt_train_classification_transformer_model.py new file mode 100644 index 0000000..a9374c1 --- /dev/null +++ b/modules/pytorch_wrapper/pt_train_classification_transformer_model.py @@ -0,0 +1,231 @@ +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 + + +class PtTrainClassificationTransformerModel: + """ + Trains 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. + loss (PTLOSS): The loss class. + epochs (int): Number of training epochs. + 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). + scheduler (PTLRSCHEDULER): Learning rate scheduler. + + 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":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, + use_gpu: bool, + early_stopping: bool, + early_stopping_rounds: int, + output_best_val_model:bool, + classification_metrics:bool, + scheduler: Optional[Union[torch.optim.lr_scheduler._LRScheduler, torch.optim.lr_scheduler.ReduceLROnPlateau]] = None, + 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 (torch.optim.Optimizer): Optimizer used for training. + loss (torch.nn.Module): The loss class. + epochs (int): Number of training epochs. + 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). + scheduler (Optional[Union[torch.optim.lr_scheduler._LRScheduler, torch.optim.lr_scheduler.ReduceLROnPlateau]]): Learning rate scheduler. + + Returns: + Tuple[torch.nn.Module, torch.Tensor, torch.Tensor]: A tuple containing the trained model, + training loss history, and validation loss history. + """ + if (scheduler is not None and + isinstance(scheduler, torch.optim.lr_scheduler.ReduceLROnPlateau) and + val_loader is None): + raise ValueError("Validation data loader needs to be specified when ReduceLROnPlateau scheduler is specified.") + + 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 + + all_preds = [] + all_labels = [] + + 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) + + 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(): + 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 = (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 scheduler: + if isinstance(scheduler, torch.optim.lr_scheduler.ReduceLROnPlateau): + scheduler.step(epoch_loss) + else: + scheduler.step() + + 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)) \ No newline at end of file diff --git a/modules/pytorch_wrapper/ptn_embedding_transformer_linear.py b/modules/pytorch_wrapper/ptn_embedding_transformer_linear.py new file mode 100644 index 0000000..80298ea --- /dev/null +++ b/modules/pytorch_wrapper/ptn_embedding_transformer_linear.py @@ -0,0 +1,251 @@ +from typing import Any, Dict, Tuple, Union, Callable, Optional +import torch +import torch.nn as nn +import torch +import torch.nn as nn +from typing import Tuple, Union, Callable + + +class EmbeddingTransformerLinear(nn.Module): + """ + A Transformer encoder-based classification model with a linear output layer. + + This module uses a standard PyTorch TransformerEncoder, followed by masked mean pooling + and a linear layer applied to the pooled representation. + + Args: + num_layers (int): Number of encoder layers. + vocabulary_size (int): Vocabulary size of the tokens. + hidden_size (int): Hidden size (d_model) of the Transformer. + nhead (int): Number of attention heads. + dim_feedforward (int): Feedforward dimension. + dropout (float): Dropout rate. + nonlinearity (str): Activation function ('relu' or 'gelu'). + layer_norm_eps (float): Epsilon for layer normalization. + batch_first (bool): Input shape format: True for (B, Seq, Feature), False for (Seq, B, Feature). + norm_first (bool): Whether normalization comes before sublayers. + bias (bool): Use bias in Transformer layers. + max_length (int): The maximum sequence length used for allocating the positional embedding. This value is not passed to the Transformer itself. + linear_output_size (int): Size of linear layer output. + linear_bias (bool): Use bias in linear layer. + + pragma: skip_doc + """ + + def __init__(self, + num_encoder_layers: int, + vocabulary_size: int, + d_model: int, + nhead: int, + dim_feedforward: int, + dropout: float, + activation: Union[str, Callable[[torch.Tensor], torch.Tensor]], + layer_norm_eps: float, + batch_first: bool, + norm_first: bool, + bias: bool, + max_length: int, + linear_output_size: int, + linear_bias: bool + ): + super().__init__() + self.batch_first = batch_first + + self.word_embedding = nn.Embedding(vocabulary_size, d_model) + self.positional_embedding = nn.Embedding(max_length, d_model) + + # FIX: Allocate position_ids buffer once to avoid reallocation in each forward pass + # self.register_buffer("position_ids_buffer", torch.arange(max_length)) # (max_length,) + + # FIX: use preallocated position_ids_buffer + # position_ids = self.position_ids_buffer[:seq_len].unsqueeze(0).expand(batch_size, -1) # (B, T) + self.position_ids = torch.arange(max_length) + + encoder_layer = nn.TransformerEncoderLayer( + d_model=d_model, + nhead=nhead, + dim_feedforward=dim_feedforward, + dropout=dropout, + activation=activation, + layer_norm_eps=layer_norm_eps, + batch_first=batch_first, + norm_first=norm_first, + bias=bias + ) + + self.encoder = nn.TransformerEncoder( + encoder_layer=encoder_layer, + num_layers=num_encoder_layers + ) + + self.linear_model = nn.Linear( + in_features=d_model, + out_features=linear_output_size, + bias=linear_bias + ) + torch.nn.init.xavier_uniform_(self.linear_model.weight) + + def forward(self, inputs: Tuple[torch.Tensor, torch.Tensor]) -> torch.Tensor: + """ + Forward pass through the Transformer encoder and linear layer. + + Args: + inputs (Tuple[torch.Tensor, torch.Tensor]): + tokens: (B, T) token IDs, + masks: (B, T) attention mask (1 = keep, 0 = pad) + + Returns: + torch.Tensor: Output of the linear layer (B, linear_output_size) + """ + tokens, masks = inputs # (B, Tokens), (B, Tokens) + if tokens.dim() != 2: + raise ValueError(f"tokens with rank {tokens.dim()} was passed.") + if masks.dim() != 2: + raise ValueError(f"masks with rank {tokens.dim()} was passed.") + if tokens.size() != masks.size(): + raise ValueError(f"tokens size masks size do not match.") + + batch_size, seq_len = tokens.size() + token_embeds = self.word_embedding(tokens) + + position_ids_repeated = self.position_ids[:seq_len].repeat(batch_size, 1).to(token_embeds.device) + positional_embeds = self.positional_embedding(position_ids_repeated) + + if token_embeds.size() != positional_embeds.size(): + raise ValueError("token_embeds & positional_embeds size mismatch.") + x = token_embeds + positional_embeds # (B, Seq, Feature) + + if not self.batch_first: + x = x.transpose(0, 1) # (Seq, B, Feature) + + x = self.encoder( + x, + src_key_padding_mask=~masks.bool() # (B, Seq) + ) + + if not self.batch_first: + x = x.transpose(0, 1) # (B, Seq, Feature) + + # FIX: More memory-efficient masked mean pooling + mask = masks.float() # (B, Seq) + x = (x * mask.unsqueeze(-1)).sum(dim=1) / mask.sum(dim=1, keepdim=True).clamp(min=1e-9) # (B, Feature) + + outputs = self.linear_model(x) # (B, linear_output_size) + return outputs + + +class PtnEmbeddingTransformerLinear: + """ + A Transformer encoder model with a linear head. + + Args: + num_layers (int): Number of encoder layers in Transformer. + vocabulary_size (int): The vocabulary size of the tokens. + hidden_size (int): Hidden dimension (d_model) for each token of the Transformer. + nhead (int): Number of attention heads. + dim_feedforward (int): Feedforward dimension. + dropout (float): Dropout rate for Transformer layers. + nonlinearity (str): Activation function ('relu' or 'gelu') used in Transformer layers. + layer_norm_eps (float): Epsilon value for layer normalization. + batch_first (bool): Input shape format: True for (B, Seq, Feature), False for (Seq, B, Feature). + norm_first (bool): Whether to apply normalization before attention/feedforward. + bias (bool): Use bias in the Transformer layers. + max_length (int): Maximum input sequence length (for positional embedding). + linear_output_size (int): Output size of the final linear layer. + linear_bias (bool): Use bias in the linear layer. + + 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": { + "num_layers": ("INT", {"default": 6, "min": 1, "max": 256}), + "vocabulary_size": ("INT", {"default": 10000, "min": 1, "max": 1e6}), + "hidden_size": ("INT", {"default": 512, "min": 1, "max": 1e6}), + "nhead": ("INT", {"default": 8, "min": 1, "max": 32}), + "dim_feedforward": ("INT", {"default": 2048, "min": 1, "max": 1e6}), + "dropout": ("FLOAT", {"default": 0.1, "min": 0, "max": 1}), + "nonlinearity": (("gelu", "relu"),), + "layer_norm_eps": ("FLOAT", {"default": 1e-5, "min": 1e-9, "max": 1e-1, "step": 1e-9}), + "batch_first": ("BOOLEAN", {"default": True}), + "norm_first": ("BOOLEAN", {"default": False}), + "bias": ("BOOLEAN", {"default": True}), + "max_length": ("INT", {"default": 512, "min": 1, "max": 1e6}), + "linear_output_size": ("INT", {"default": 1, "min": 1, "max": 1e6}), + "linear_bias": ("BOOLEAN", {"default": True}), + } + } + + RETURN_TYPES: Tuple[str] = ("PTMODEL",) + FUNCTION: str = "f" + CATEGORY: str = "Training" + + @classmethod + def IS_CHANGED(cls, **kw): + return float("NaN") + + def f(self, + num_layers: int, + vocabulary_size: int, + hidden_size: int, + nhead: int, + dim_feedforward: int, + dropout: float, + nonlinearity: str, + layer_norm_eps: float, + batch_first: bool, + norm_first: bool, + bias: bool, + max_length: int, + linear_output_size: int, + linear_bias: bool + ) -> Tuple[nn.Module]: + """ + Constructs a Transformer model with a linear output layer. + + Args: + num_layers (int): Number of encoder layers. + vocabulary_size (int): Vocabulary size of the tokens. + hidden_size (int): Hidden size (d_model) of the Transformer. + nhead (int): Number of attention heads. + dim_feedforward (int): Feedforward dimension. + dropout (float): Dropout rate. + nonlinearity (str): Activation function ('relu' or 'gelu'). + layer_norm_eps (float): Epsilon for layer normalization. + batch_first (bool): Input shape format: True for (B, Seq, Feature), False for (Seq, B, Feature). + norm_first (bool): Whether normalization comes before sublayers. + bias (bool): Use bias in Transformer layers. + max_length (int): Maximum sequence length. + linear_output_size (int): Size of linear layer output. + linear_bias (bool): Use bias in linear layer. + + Returns: + Tuple[nn.Module]: A tuple containing the instantiated PyTorch model. + """ + with torch.inference_mode(False): + model = EmbeddingTransformerLinear( + num_encoder_layers=num_layers, + vocabulary_size=vocabulary_size, + d_model=hidden_size, + nhead=nhead, + dim_feedforward=dim_feedforward, + dropout=dropout, + activation=nonlinearity, + layer_norm_eps=layer_norm_eps, + batch_first=batch_first, + norm_first=norm_first, + bias=bias, + max_length=max_length, + linear_output_size=linear_output_size, + linear_bias=linear_bias + ) + return (model,) + diff --git a/modules/pytorch_wrapper/pto_lr_scheduler_cosine_annealing.py b/modules/pytorch_wrapper/pto_lr_scheduler_cosine_annealing.py index c3c58ed..51a30e6 100644 --- a/modules/pytorch_wrapper/pto_lr_scheduler_cosine_annealing.py +++ b/modules/pytorch_wrapper/pto_lr_scheduler_cosine_annealing.py @@ -54,7 +54,7 @@ class PtoLrSchedulerCosineAnnealing: "required": { "optimizer": ("PTOPTIMIZER", {}), "num_epochs": ("INT", {"default": 10, "min": 1, "max": int(1e8)}), - "minimum_lr": ("FLOAT", {"default": 0, "min": 0, "max": 1.0, "step": 1e-8}), + "minimum_lr": ("FLOAT", {"default": 0, "min": 0, "max": 1.0, "step": 1e-9}), } } diff --git a/modules/pytorch_wrapper/ptv_hf_dataset_with_token_encode.py b/modules/pytorch_wrapper/ptv_hf_dataset_with_token_encode.py index 3183ba2..d75b9f1 100644 --- a/modules/pytorch_wrapper/ptv_hf_dataset_with_token_encode.py +++ b/modules/pytorch_wrapper/ptv_hf_dataset_with_token_encode.py @@ -12,6 +12,7 @@ class PtvHfDatasetWithTokenEncode: sample_field_name (str): Field name for text samples. 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. category: PyTorch wrapper - Training """ @@ -31,6 +32,7 @@ class PtvHfDatasetWithTokenEncode: "sample_field_name": ("STRING", {"default": "text","multiline": False}), "label_field_name": ("STRING", {"default": "label","multiline": False}), "encode": ("PTCALLABLE", {}), + "remove_html_tags": ("BOOLEAN", {"default": False}), } } @@ -47,7 +49,8 @@ class PtvHfDatasetWithTokenEncode: split: str, sample_field_name: str, label_field_name: str, - encode: Callable) -> Tuple: + encode: Callable, + remove_html_tags:bool) -> Tuple: """ Loads a dataset from Hugging Face with specified parameters. @@ -57,6 +60,7 @@ class PtvHfDatasetWithTokenEncode: sample_field_name (str): Field name for text samples. 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. Returns: Tuple: A tuple containing the dataset instance. @@ -71,5 +75,6 @@ class PtvHfDatasetWithTokenEncode: split=split, sample_field_name=sample_field_name, label_field_name=label_field_name, - encode=encode) + encode=encode, + remove_html_tags=remove_html_tags) return (dc,) diff --git a/pyproject.toml b/pyproject.toml index e70927d..7fc217a 100644 --- a/pyproject.toml +++ b/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.3.4" +version = "3.4.0" license = {file = "LICENSE"} [project.urls] diff --git a/tests/pytorch_wrapper/hf_dataset_with_tokenize_test.py b/tests/pytorch_wrapper/hf_dataset_with_tokenize_test.py index 51d9277..160dd20 100644 --- a/tests/pytorch_wrapper/hf_dataset_with_tokenize_test.py +++ b/tests/pytorch_wrapper/hf_dataset_with_tokenize_test.py @@ -37,6 +37,13 @@ class TestHfDatasetWithTokenEncode(unittest.TestCase): "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) @@ -58,5 +65,25 @@ class TestHfDatasetWithTokenEncode(unittest.TestCase): 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() diff --git a/tests/pytorch_wrapper/pt_train_embedding_transformer_linear_text_classification_model_test.py b/tests/pytorch_wrapper/pt_train_embedding_transformer_linear_text_classification_model_test.py new file mode 100644 index 0000000..610f0bc --- /dev/null +++ b/tests/pytorch_wrapper/pt_train_embedding_transformer_linear_text_classification_model_test.py @@ -0,0 +1,139 @@ +""" +Dataset URL:https://huggingface.co/datasets/stanfordnlp/imdb +""" +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_embedding_transformer_linear import PtnEmbeddingTransformerLinear +from pytorch_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 +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 TestPtTrainEmbeddingTransformerLinearTextClassificationModel(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.load_node = SpLoadModel() + self.sp = self.load_node.f("spiece.model")[0] + self.encode_node = SpEncode() + self.max_length=512 + encode = self.encode_node.f( + spmodel=self.sp, + padding=True, + padding_method="max_length", + padding_value=0, + truncation=True, + max_length=self.max_length + )[0] + + train_dataset = HfDatasetWithTokenEncode( + "imdb", # dataset_name + "train", #split + encode=encode, + remove_html_tags=True + ) + + test_dataset = HfDatasetWithTokenEncode( + "imdb", # dataset_name + "test", #split + encode=encode, + remove_html_tags=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 = "transformer_text_classification_test.pt" + + self.model_node = PtnEmbeddingTransformerLinear() + self.model = self.model_node.f( + 2, # num_encoder + 32000, # vocabulary_size + 256, # hidden_size + 8, # nhead + 512, # dim_feedforward + 0.5, # dropout + "gelu", # activation + 1e-5, # layer_norm_eps, + True, # batch_first + False, # norm_first + True, # bias + self.max_length, # max_length of seq + 1, # linear_output_size + True # linear_bias + )[0] + + self.optimizer = PtoAdamW().f(self.model, 3e-4, 0.9, 0.999, + 0.01, # weight decay + False # amsgrad + )[0] + + self.minimum_lr = 1e-9 + self.epochs = 10 + 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() diff --git a/tests/pytorch_wrapper/ptn_embedding_transformer_linear_test.py b/tests/pytorch_wrapper/ptn_embedding_transformer_linear_test.py new file mode 100644 index 0000000..8045578 --- /dev/null +++ b/tests/pytorch_wrapper/ptn_embedding_transformer_linear_test.py @@ -0,0 +1,43 @@ +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_embedding_transformer_linear import PtnEmbeddingTransformerLinear + + +class TestPtnEmbeddingTransformerLinear(unittest.TestCase): + def setUp(self): + """Set up test instance.""" + self.model_node = PtnEmbeddingTransformerLinear() + + def test_1(self): + self.model = self.model_node.f( + 1, # num_encoder + 32000, # vocabulary_size + 32, # hidden_size + 8, # nhead + 64, # dim_feedforward + 0.1, # dropout + "gelu", # activation + 1e-5, # layer_norm_eps, + True, # batch_first + False, # norm_first + True, # bias + 512, # max_length of seq + 1, # linear_output_size + True # linear_bias + )[0] + + x = torch.randint(0, 32000, (8 * 77,)).to(torch.int64) + x = x.view(8, 77) + masks = torch.ones_like(x) + outputs = self.model((x, masks)) + self.assertEqual(outputs.size(), torch.Size([8, 1])) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/pytorch_wrapper/ptv_hf_dataset_with_token_encode_test.py b/tests/pytorch_wrapper/ptv_hf_dataset_with_token_encode_test.py index f674719..a01d43e 100644 --- a/tests/pytorch_wrapper/ptv_hf_dataset_with_token_encode_test.py +++ b/tests/pytorch_wrapper/ptv_hf_dataset_with_token_encode_test.py @@ -36,7 +36,8 @@ class TestPtvHfDatasetWithTokenEncode(unittest.TestCase): "train", #split "text", "label", - encode + encode, + False # remove_html_tags )[0]