, 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]