From faf3a8f0e510c74416bd95fb1b483291ac4ef40d Mon Sep 17 00:00:00 2001 From: HowToSD Date: Thu, 27 Mar 2025 11:51:21 -0700 Subject: [PATCH] Added Embedding layer and Embedding RNN Layer node. --- __init__.py | 18 ++ credit/credit.md | 254 ++++++++++++++++++ docs/reference/node_reference.md | 6 + docs/reference/pt_tokenizer.md | 26 ++ docs/reference/ptn_embedding.md | 24 ++ docs/reference/ptn_embedding_rnn_linear.md | 40 +++ .../ptv_hf_dataset_with_token_encode.md | 28 ++ docs/reference/ptv_hf_glove_dataset.md | 2 +- docs/reference/sp_encode.md | 36 +++ docs/reference/sp_load_model.md | 17 ++ docs/training_rnn_for_classification.md | 19 +- .../embedding_rnn_classification.json | 1 + .../hf_dataset_with_token_encode.py | 62 +++++ modules/pytorch_wrapper/pt_tokenizer.py | 49 ++++ modules/pytorch_wrapper/pt_train_rnn_model.py | 10 +- modules/pytorch_wrapper/ptn_embedding.py | 65 +++++ .../ptn_embedding_rnn_linear.py | 197 ++++++++++++++ .../ptv_hf_dataset_with_token_encode.py | 75 ++++++ .../pytorch_wrapper/ptv_hf_glove_dataset.py | 4 +- modules/pytorch_wrapper/utils.py | 94 ++++++- modules/sentencepiece_wrapper/sp_encode.py | 134 +++++++++ .../sentencepiece_wrapper/sp_load_model.py | 47 ++++ pyproject.toml | 2 +- requirements.txt | 1 + .../hf_dataset_with_tokenize_test.py | 62 +++++ tests/pytorch_wrapper/pt_tokenizer_test.py | 43 +++ ...n_linear_text_classification_model_test.py | 133 +++++++++ .../ptn_embedding_rnn_linear_test.py | 41 +++ tests/pytorch_wrapper/ptn_embedding_test.py | 61 +++++ .../ptv_hf_dataset_with_token_encode_test.py | 65 +++++ .../utils_pad_truncate_sequence_test.py | 120 +++++++++ tests/sentencepiece_wrapper/sp_encode_test.py | 146 ++++++++++ .../sp_load_model_test.py | 29 ++ 33 files changed, 1899 insertions(+), 12 deletions(-) create mode 100644 docs/reference/pt_tokenizer.md create mode 100644 docs/reference/ptn_embedding.md create mode 100644 docs/reference/ptn_embedding_rnn_linear.md create mode 100644 docs/reference/ptv_hf_dataset_with_token_encode.md create mode 100644 docs/reference/sp_encode.md create mode 100644 docs/reference/sp_load_model.md create mode 100644 examples/workflows/embedding_rnn_classification.json create mode 100644 modules/pytorch_wrapper/hf_dataset_with_token_encode.py create mode 100644 modules/pytorch_wrapper/pt_tokenizer.py create mode 100644 modules/pytorch_wrapper/ptn_embedding.py create mode 100644 modules/pytorch_wrapper/ptn_embedding_rnn_linear.py create mode 100644 modules/pytorch_wrapper/ptv_hf_dataset_with_token_encode.py create mode 100644 modules/sentencepiece_wrapper/sp_encode.py create mode 100644 modules/sentencepiece_wrapper/sp_load_model.py create mode 100644 tests/pytorch_wrapper/hf_dataset_with_tokenize_test.py create mode 100644 tests/pytorch_wrapper/pt_tokenizer_test.py create mode 100644 tests/pytorch_wrapper/pt_train_embedding_rnn_linear_text_classification_model_test.py create mode 100644 tests/pytorch_wrapper/ptn_embedding_rnn_linear_test.py create mode 100644 tests/pytorch_wrapper/ptn_embedding_test.py create mode 100644 tests/pytorch_wrapper/ptv_hf_dataset_with_token_encode_test.py create mode 100644 tests/pytorch_wrapper/utils_pad_truncate_sequence_test.py create mode 100644 tests/sentencepiece_wrapper/sp_encode_test.py create mode 100644 tests/sentencepiece_wrapper/sp_load_model_test.py diff --git a/__init__.py b/__init__.py index 03dcef8..4d5ae3f 100644 --- a/__init__.py +++ b/__init__.py @@ -107,6 +107,7 @@ from .modules.pytorch_wrapper.pt_to_latent import PtToLatent from .modules.pytorch_wrapper.pt_to_numpy import PtToNumpy from .modules.pytorch_wrapper.pt_to_rgb_tensors import PtToRgbTensors 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_model import PtTrainModel @@ -156,6 +157,8 @@ from .modules.pytorch_wrapper.ptn_chained_model import PtnChainedModel from .modules.pytorch_wrapper.ptn_conv2d import PtnConv2d 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_gru import PtnGRU from .modules.pytorch_wrapper.ptn_gru_linear import PtnGRULinear from .modules.pytorch_wrapper.ptn_huber_loss import PtnHuberLoss @@ -186,12 +189,15 @@ from .modules.pytorch_wrapper.pto_simple import PtoSimple from .modules.pytorch_wrapper.ptv_dataset import PtvDataset from .modules.pytorch_wrapper.ptv_dataset_len import PtvDatasetLen from .modules.pytorch_wrapper.ptv_dataset_loader import PtvDatasetLoader +from .modules.pytorch_wrapper.ptv_hf_dataset_with_token_encode import PtvHfDatasetWithTokenEncode from .modules.pytorch_wrapper.ptv_hf_glove_dataset import PtvHfGloveDataset from .modules.pytorch_wrapper.ptv_image_folder_dataset import PtvImageFolderDataset from .modules.pytorch_wrapper.ptv_sequential_tensor_dataset import PtvSequentialTensorDataset from .modules.pytorch_wrapper.ptv_transforms_data_augment import PtvTransformsDataAugment from .modules.pytorch_wrapper.ptv_transforms_resize import PtvTransformsResize from .modules.pytorch_wrapper.ptv_transforms_to_tensor import PtvTransformsToTensor +from .modules.sentencepiece_wrapper.sp_encode import SpEncode +from .modules.sentencepiece_wrapper.sp_load_model import SpLoadModel T = TypeVar("T") @@ -308,6 +314,7 @@ NODE_CLASS_MAPPINGS: Dict[str, Type[T]] = { "PtToNumpy": PtToNumpy, "PtToRgbTensors": PtToRgbTensors, "PtToUint8": PtToUint8, + "PtTokenizer": PtTokenizer, "PtTrainClassificationModel": PtTrainClassificationModel, "PtTrainClassificationModelLr": PtTrainClassificationModelLr, "PtTrainModel": PtTrainModel, @@ -357,6 +364,8 @@ NODE_CLASS_MAPPINGS: Dict[str, Type[T]] = { "PtnConv2d": PtnConv2d, "PtnConvModel": PtnConvModel, "PtnCrossEntropyLoss": PtnCrossEntropyLoss, + "PtnEmbedding": PtnEmbedding, + "PtnEmbeddingRNNLinear": PtnEmbeddingRNNLinear, "PtnGRU": PtnGRU, "PtnGRULinear": PtnGRULinear, "PtnHuberLoss": PtnHuberLoss, @@ -387,12 +396,15 @@ NODE_CLASS_MAPPINGS: Dict[str, Type[T]] = { "PtvDataset": PtvDataset, "PtvDatasetLen": PtvDatasetLen, "PtvDatasetLoader": PtvDatasetLoader, + "PtvHfDatasetWithTokenEncode": PtvHfDatasetWithTokenEncode, "PtvHfGloveDataset": PtvHfGloveDataset, "PtvImageFolderDataset": PtvImageFolderDataset, "PtvSequentialTensorDataset": PtvSequentialTensorDataset, "PtvTransformsDataAugment": PtvTransformsDataAugment, "PtvTransformsResize": PtvTransformsResize, "PtvTransformsToTensor": PtvTransformsToTensor, + "SpEncode": SpEncode, + "SpLoadModel": SpLoadModel, } @@ -509,6 +521,7 @@ NODE_DISPLAY_NAME_MAPPINGS: Dict[str, str] = { "PtToNumpy": "Pt To Numpy", "PtToRgbTensors": "Pt To Rgb Tensors", "PtToUint8": "Pt To Uint8", + "PtTokenizer": "Pt Tokenizer", "PtTrainClassificationModel": "Pt Train Classification Model", "PtTrainClassificationModelLr": "Pt Train Classification Model Lr", "PtTrainModel": "Pt Train Model", @@ -558,6 +571,8 @@ NODE_DISPLAY_NAME_MAPPINGS: Dict[str, str] = { "PtnConv2d": "Ptn Conv 2d", "PtnConvModel": "Ptn Conv Model", "PtnCrossEntropyLoss": "Ptn Cross Entropy Loss", + "PtnEmbedding": "Ptn Embedding", + "PtnEmbeddingRNNLinear": "Ptn Embedding RNN Linear", "PtnGRU": "Ptn GRU", "PtnGRULinear": "Ptn GRU Linear", "PtnHuberLoss": "Ptn Huber Loss", @@ -588,12 +603,15 @@ NODE_DISPLAY_NAME_MAPPINGS: Dict[str, str] = { "PtvDataset": "Ptv Dataset", "PtvDatasetLen": "Ptv Dataset Len", "PtvDatasetLoader": "Ptv Dataset Loader", + "PtvHfDatasetWithTokenEncode": "Ptv Hf Dataset With Token Encode", "PtvHfGloveDataset": "Ptv Hf Glove Dataset", "PtvImageFolderDataset": "Ptv Image Folder Dataset", "PtvSequentialTensorDataset": "Ptv Sequential Tensor Dataset", "PtvTransformsDataAugment": "Ptv Transforms Data Augment", "PtvTransformsResize": "Ptv Transforms Resize", "PtvTransformsToTensor": "Ptv Transforms To Tensor", + "SpEncode": "Sp Encode", + "SpLoadModel": "Sp Load Model", } diff --git a/credit/credit.md b/credit/credit.md index 8b416b7..cac3c00 100644 --- a/credit/credit.md +++ b/credit/credit.md @@ -28,3 +28,257 @@ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. ``` + +## spiece.model + +This tokenizer model (`spiece.model`) was downloaded from: + +https://huggingface.co/google-t5/t5-base/tree/main + +It is part of the T5 model published by Google and is licensed under the Apache License 2.0. + +© Google Research. Licensed under the Apache License, Version 2.0. +You may obtain a copy of the license at: + +http://www.apache.org/licenses/LICENSE-2.0 + +A copy of the full license text is included below. +``` +--- +title: Apache License 2.0 +spdx-id: Apache-2.0 +redirect_from: /licenses/apache/ +featured: true +hidden: false + +description: A permissive license whose main conditions require preservation of copyright and license notices. Contributors provide an express grant of patent rights. Licensed works, modifications, and larger works may be distributed under different terms and without source code. + +how: Create a text file (typically named LICENSE or LICENSE.txt) in the root of your source code and copy the text of the license into the file. + +note: The Apache Software Foundation recommends taking the additional step of adding a boilerplate notice to the header of each source file. You can find the notice in the appendix at the very end of the license text. + +using: + Kubernetes: https://github.com/kubernetes/kubernetes/blob/master/LICENSE + PDF.js: https://github.com/mozilla/pdf.js/blob/master/LICENSE + Swift: https://github.com/apple/swift/blob/main/LICENSE.txt + +permissions: + - commercial-use + - modifications + - distribution + - patent-use + - private-use + +conditions: + - include-copyright + - document-changes + +limitations: + - trademark-use + - liability + - warranty + +--- + + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. +``` \ No newline at end of file diff --git a/docs/reference/node_reference.md b/docs/reference/node_reference.md index cb276b0..bcc50f0 100644 --- a/docs/reference/node_reference.md +++ b/docs/reference/node_reference.md @@ -30,6 +30,7 @@ ## PyTorch wrapper - Callable | Node | Description | | --- | --- | +| [Pt Tokenizer](pt_tokenizer.md) | The tokenizer to encode a string or a list of string to token IDs. | | [Ptf GELU](ptf_gelu.md) | The GELU activation function. | | [Ptf LeakyReLU](ptf_leaky_relu.md) | The LeakyReLU activation function. | | [Ptf Log Softmax](ptf_log_softmax.md) | The log softmax activation function. | @@ -155,6 +156,7 @@ | [Ptn Chained Model](ptn_chained_model.md) | Constructs a chained PyTorch model. | | [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 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. | @@ -167,6 +169,7 @@ | [Ptn Pre Add Channel Axis](ptn_pre_add_channel_axis.md) | Adds a channel axis after the batch axis if the input is rank 3 (bs, h, w) | | [Ptn Pre Flatten](ptn_pre_flatten.md) | Flattens the input tensor before processing the tensor in the specified model. | | [Ptn RNN](ptn_rnn.md) | A recurrent neural network (RNN) model consisting of one or more of a recurrent layer. | +| [Ptn RNN Linear](ptn_embedding_rnn_linear.md) | A recurrent neural network (RNN) model with a linear head. | | [Ptn RNN Linear](ptn_rnn_linear.md) | A recurrent neural network (RNN) model with a linear head. | | [Ptn Resnet Model](ptn_resnet_model.md) | A Resnet model consisting of multiple Resnet layers. | ## PyTorch wrapper - Optimizer @@ -240,12 +243,15 @@ | [Ptv Dataset](ptv_dataset.md) | A Torchvision Dataset class wrapper. | | [Ptv Dataset Len](ptv_dataset_len.md) | A wrapper class that calls Python len on a dataset. | | [Ptv Dataset Loader](ptv_dataset_loader.md) | A node to combine the dataset and data loader into a single node. | +| [Ptv Hf Dataset With Token Encode](ptv_hf_dataset_with_token_encode.md) | A PyTorch Dataset class wrapper for a Hugging Face dataset that converts text to token IDs. | | [Ptv Hf Glove Dataset](ptv_hf_glove_dataset.md) | A PyTorch Dataset class wrapper for a Hugging Face dataset that converts text to embedding using Glove. | | [Ptv Image Folder Dataset](ptv_image_folder_dataset.md) | A Torchvision ImageFolder Dataset class wrapper. | | [Ptv Sequential Tensor Dataset](ptv_sequential_tensor_dataset.md) | Creates a sequential tensor Dataset. | | [Ptv Transforms Data Augment](ptv_transforms_data_augment.md) | Applies data augmentation transformations to dataset elements. | | [Ptv Transforms Resize](ptv_transforms_resize.md) | Resizes and transforms elements of dataset to PyTorch tensors. | | [Ptv Transforms To Tensor](ptv_transforms_to_tensor.md) | Transforms elements of dataset to PyTorch tensors. | +| [Sp Encode](sp_encode.md) | SentencePiece wrapper for encoding text into token ID tensors. | +| [Sp Load Model](sp_load_model.md) | A wrapper class for loading a SentencePiece tokenization model. | ## PyTorch wrapper - Transform | Node | Description | | --- | --- | diff --git a/docs/reference/pt_tokenizer.md b/docs/reference/pt_tokenizer.md new file mode 100644 index 0000000..6a6b742 --- /dev/null +++ b/docs/reference/pt_tokenizer.md @@ -0,0 +1,26 @@ +# Pt Tokenizer +The tokenizer to encode a string or a list of string to token IDs. + + Args: + encode (PTCALLABLE): The reference to a token encoder function. + text_list (PYLIST): The python list of strings to tokenize. + + Returns: + A tuple containing the token IDs and the mask in torch.Tensor. + +## Input +| Name | Data type | +|---|---| +| encode | Ptcallable | +| text_list | Pylist | + +## Output +| Data type | +|---| +| Tensor | +| Tensor | + +
+Category: PyTorch wrapper - Callable + +ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved. diff --git a/docs/reference/ptn_embedding.md b/docs/reference/ptn_embedding.md new file mode 100644 index 0000000..a14b2ac --- /dev/null +++ b/docs/reference/ptn_embedding.md @@ -0,0 +1,24 @@ +# Ptn Embedding +Constructs an embedding layer. + + Args: + num_embeddings (int): Vocabulary size. + embedding_dim (int): The size of each embedding vector. + padding_idx (int): Token ID to treat as padding. If a non-negative integer is specified, the embedding at that index will be excluded from gradient updates. Set to -1 to disable padding behavior (default). + +## Input +| Name | Data type | +|---|---| +| num_embeddings | Int | +| embedding_size | Int | +| padding_idx | Int | + +## 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/ptn_embedding_rnn_linear.md b/docs/reference/ptn_embedding_rnn_linear.md new file mode 100644 index 0000000..6ac0472 --- /dev/null +++ b/docs/reference/ptn_embedding_rnn_linear.md @@ -0,0 +1,40 @@ +# Ptn RNN Linear +A recurrent neural network (RNN) model with a linear head. + + Args: + vocabulary_size (int): The vocabulary size of the tokens. + input_size (int): The number of input features. + hidden_size (int): The number of output features of the hidden layer matrix. + num_layers (int): Number of hidden layers in RNN. + nonlinearity (str): Activation function to apply after each hidden layer. Specify 'tanh' or 'relu'. + bias (bool): Use bias or not. + batch_first (bool): Indicates the shape of the input. Input is assumed to be a rank-3 tensor. If True, the shape is [Batch, Seq, Token]. If False, the shape is [Seq, Batch, Token]. Note that default is set to `True` unlike PyTorch's RNN default model parameter. + dropout (float): Dropout probability to be applied to the output of intermediate layers if non-zero. This will not be applied to the last layer's output. + bidirectional (bool): Processes the input in a backward direction and append the result to the output of the forward direction. + linear_output_size (int): The number of output features of the linear layer. + linar_bias (bool): Use bias or not in linear layer. + +## Input +| Name | Data type | +|---|---| +| vocabulary_size | Int | +| input_size | Int | +| hidden_size | Int | +| num_layers | Int | +| nonlinearity | | +| bias | Boolean | +| batch_first | Boolean | +| dropout | Float | +| bidirectional | Boolean | +| 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 new file mode 100644 index 0000000..f1f27af --- /dev/null +++ b/docs/reference/ptv_hf_dataset_with_token_encode.md @@ -0,0 +1,28 @@ +# Ptv Hf Dataset With Token Encode +A PyTorch Dataset class wrapper for a Hugging Face dataset that converts text to token IDs. + + Args: + name (str): Name of the dataset class. + split (str): The name of the subset of dataset to retrieve or the split specifier. + 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. + +## Input +| Name | Data type | +|---|---| +| name | String | +| split | String | +| sample_field_name | String | +| label_field_name | String | +| encode | Ptcallable | + +## Output +| Data type | +|---| +| Ptvdataset | + +
+Category: PyTorch wrapper - Training + +ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved. diff --git a/docs/reference/ptv_hf_glove_dataset.md b/docs/reference/ptv_hf_glove_dataset.md index 812ffdd..7f44487 100644 --- a/docs/reference/ptv_hf_glove_dataset.md +++ b/docs/reference/ptv_hf_glove_dataset.md @@ -2,7 +2,7 @@ A PyTorch Dataset class wrapper for a Hugging Face dataset that converts text to embedding using Glove. Args: - name (str): Name of the dataset class (e.g., MNIST, FashionMNIST). + name (str): Name of the dataset class. split (str): The name of the subset of dataset to retrieve or the split specifier. embed_dim (int): Embedding dimension for Glove. max_seq_len (int): Sequence length. diff --git a/docs/reference/sp_encode.md b/docs/reference/sp_encode.md new file mode 100644 index 0000000..f180561 --- /dev/null +++ b/docs/reference/sp_encode.md @@ -0,0 +1,36 @@ +# Sp Encode +SentencePiece wrapper for encoding text into token ID tensors. + + Args: + spmodel (SentencePieceProcessor): Trained SentencePiece 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`. + padding_value (Union[int, torch.Tensor]): A scalar value or scalar tensor used for padding. + truncation (bool): Whether to truncate sequences to `max_length`. + max_length (int): The target length for truncation or padding. + + 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 | +|---|---| +| spmodel | Spmodel | +| padding | Boolean | +| padding_method | | +| padding_value | Int | +| truncation | Boolean | +| max_length | Int | + +## Output +| Data type | +|---| +| Ptcallable | + +
+Category: PyTorch wrapper - Training + +ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved. diff --git a/docs/reference/sp_load_model.md b/docs/reference/sp_load_model.md new file mode 100644 index 0000000..3302325 --- /dev/null +++ b/docs/reference/sp_load_model.md @@ -0,0 +1,17 @@ +# Sp Load Model +A wrapper class for loading a SentencePiece tokenization model. + +## Input +| Name | Data type | +|---|---| +| model_path | String | + +## Output +| Data type | +|---| +| Spmodel | + +
+Category: PyTorch wrapper - Training + +ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved. diff --git a/docs/training_rnn_for_classification.md b/docs/training_rnn_for_classification.md index 4def5b8..e206e7f 100644 --- a/docs/training_rnn_for_classification.md +++ b/docs/training_rnn_for_classification.md @@ -3,7 +3,11 @@ You can use **ComfyUI-Pt-Wrapper** to train an RNN for text classification, such as sentiment analysis on a movie review dataset. You can find the workflow at `examples/workflows/rnn_classification.json`. -You can also use GRU workflow at `examples/workflows/gru_classification.json` or LSTM workflow at `examples/workflows/lstm_classification.json`. + +There are also other workflows: +* GRU workflow: `examples/workflows/gru_classification.json` +* LSTM workflow: `examples/workflows/lstm_classification.json`. +* RNN workflow with training embedding from scratch: `examples/workflows/embedding_rnn_classification.json` (Refer to the end of this document for more information) ![Workflow](images/rnn_classification.png) @@ -44,3 +48,16 @@ You can also enable `classification_metrics` to print validation accuracy to the 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. Both GRU & LSTM workflows achieve approximately **87% validation accuracy**. + +# RNN Workflow with Training Embedding from Scratch + +**Workflow:** `examples/workflows/embedding_rnn_classification.json` + +This workflow demonstrates how to use a SentencePiece tokenizer to convert text into token IDs and train an RNN classification model with an embedding layer initialized from scratch. The embedding layer learns to map token IDs to word embeddings during training. + +SentencePiece requires a pretrained model file for tokenization. A compatible model file is included in this package, sourced from [google-t5/t5-base](https://huggingface.co/google-t5/t5-base/tree/main), part of the T5 model published by Google under the Apache License 2.0. + +When configuring the workflow, set the vocabulary size to **32000** to match the included model. + +This setup achieves approximately **83% validation accuracy** on the Hugging Face version of the IMDB dataset. + diff --git a/examples/workflows/embedding_rnn_classification.json b/examples/workflows/embedding_rnn_classification.json new file mode 100644 index 0000000..f558e48 --- /dev/null +++ b/examples/workflows/embedding_rnn_classification.json @@ -0,0 +1 @@ +{"last_node_id":25,"last_link_id":28,"nodes":[{"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":[14],"slot_index":0}],"properties":{"Node name for S&R":"PtDataLoader"},"widgets_values":[128,false,"{\"num_workers\":1}"]},{"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":[13],"slot_index":0}],"properties":{"Node name for S&R":"PtDataLoader"},"widgets_values":[128,true,"{\"num_workers\":1}"]},{"id":14,"type":"PtTrainRNNModel","pos":[3111.27294921875,1735.56494140625],"size":[315,346],"flags":{},"order":10,"mode":0,"inputs":[{"name":"model","localized_name":"model","type":"PTMODEL","link":23},{"name":"train_loader","localized_name":"train_loader","type":"PTDATALOADER","link":13},{"name":"optimizer","localized_name":"optimizer","type":"PTOPTIMIZER","link":17},{"name":"loss_function","localized_name":"loss_function","type":"PTLOSS","link":18},{"name":"scheduler","localized_name":"scheduler","type":"PTLRSCHEDULER","shape":7,"link":16},{"name":"val_loader","localized_name":"val_loader","type":"PTDATALOADER","shape":7,"link":14},{"name":"h_0","localized_name":"h_0","type":"TENSOR","shape":7,"link":null}],"outputs":[{"name":"Model","localized_name":"Model","type":"PTMODEL","links":[19],"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":"PtTrainRNNModel"},"widgets_values":[true,50,true,false,10,false,true,true]},{"id":16,"type":"PtnBCEWithLogitsLoss","pos":[2668.492919921875,1800.2408447265625],"size":[315,58],"flags":{},"order":0,"mode":0,"inputs":[],"outputs":[{"name":"PTLOSS","localized_name":"PTLOSS","type":"PTLOSS","links":[18],"slot_index":0}],"properties":{"Node name for S&R":"PtnBCEWithLogitsLoss"},"widgets_values":["mean"]},{"id":9,"type":"PtoAdam","pos":[2668.492919921875,1191.617431640625],"size":[315,106],"flags":{},"order":3,"mode":0,"inputs":[{"name":"model","localized_name":"model","type":"PTMODEL","link":22}],"outputs":[{"name":"PTOPTIMIZER","localized_name":"PTOPTIMIZER","type":"PTOPTIMIZER","links":[15,17],"slot_index":0}],"properties":{"Node name for S&R":"PtoAdam"},"widgets_values":[0.001,0.9,0.999]},{"id":15,"type":"PtoLrSchedulerReduceOnPlateau","pos":[3111.27294921875,1377.6142578125],"size":[314.1559143066406,82],"flags":{},"order":5,"mode":0,"inputs":[{"name":"optimizer","localized_name":"optimizer","type":"PTOPTIMIZER","link":15}],"outputs":[{"name":"PTLRSCHEDULER","localized_name":"PTLRSCHEDULER","type":"PTLRSCHEDULER","links":[16],"slot_index":0}],"properties":{"Node name for S&R":"PtoLrSchedulerReduceOnPlateau"},"widgets_values":[3,0.5000000000000002]},{"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_rnn_classification.pt"]},{"id":20,"type":"PtnEmbeddingRNNLinear","pos":[2671.97216796875,1393.1790771484375],"size":[315,298],"flags":{},"order":1,"mode":0,"inputs":[],"outputs":[{"name":"PTMODEL","localized_name":"PTMODEL","type":"PTMODEL","links":[22,23],"slot_index":0}],"properties":{"Node name for S&R":"PtnEmbeddingRNNLinear"},"widgets_values":[32000,300,128,2,"tanh",true,true,0.1,true,1,true]},{"id":21,"type":"PtvHfDatasetWithTokenEncode","pos":[1907.6329345703125,1746.281982421875],"size":[271.5191650390625,130.8897247314453],"flags":{},"order":6,"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"]},{"id":22,"type":"SpEncode","pos":[1502.9549560546875,1913.823974609375],"size":[315,154],"flags":{},"order":4,"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,300]},{"id":25,"type":"SpLoadModel","pos":[1246.61181640625,1912.7467041015625],"size":[210,60.936126708984375],"flags":{},"order":2,"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":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":19}],"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_rnn_classification.pt"]},{"id":24,"type":"PtvHfDatasetWithTokenEncode","pos":[1911.3275146484375,2022.5567626953125],"size":[271.5191650390625,130.8897247314453],"flags":{},"order":7,"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"]}],"links":[[13,7,0,14,1,"PTDATALOADER"],[14,12,0,14,5,"PTDATALOADER"],[15,9,0,15,0,"PTOPTIMIZER"],[16,15,0,14,4,"PTLRSCHEDULER"],[17,9,0,14,2,"PTOPTIMIZER"],[18,16,0,14,3,"PTLOSS"],[19,14,0,17,0,"PTMODEL"],[20,17,0,19,0,"STRING"],[22,20,0,9,0,"PTMODEL"],[23,20,0,14,0,"PTMODEL"],[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"]],"groups":[],"config":{},"extra":{"ds":{"scale":0.9366078361504352,"offset":[-1728.9490496794083,-1464.6600418938542]},"node_versions":{"pt-wrapper":"dfcc202c5004f0fde0a527ec54c4e9724a73ec1e","data-analysis":"68dd090367adffa569c3b3993d56e422f4516d2d"},"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 new file mode 100644 index 0000000..b67e799 --- /dev/null +++ b/modules/pytorch_wrapper/hf_dataset_with_token_encode.py @@ -0,0 +1,62 @@ +from typing import Tuple, Callable +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: Callable=None) -> 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 (Callable): Token encode function + """ + 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 + + + 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] + tokens, masks = self.encode(sample[self.sample_field_name]) + + # Get a corresponding label + label = torch.tensor(sample[self.label_field_name], dtype=torch.long) + return (tokens, masks), label + + def __len__(self) -> int: + """ + Returns the number of samples in the dataset. + + Returns: + int: Dataset length. + """ + return len(self.dataset) diff --git a/modules/pytorch_wrapper/pt_tokenizer.py b/modules/pytorch_wrapper/pt_tokenizer.py new file mode 100644 index 0000000..9de67ba --- /dev/null +++ b/modules/pytorch_wrapper/pt_tokenizer.py @@ -0,0 +1,49 @@ +from typing import Any, Dict, Callable, Tuple, List, Union + + +class PtTokenizer: + """ + The tokenizer to encode a string or a list of string to token IDs. + + Args: + encode (PTCALLABLE): The reference to a token encoder function. + text_list (PYLIST): The python list of strings to tokenize. + + Returns: + A tuple containing the token IDs and the mask in torch.Tensor. + + category: PyTorch wrapper - Callable + """ + + @classmethod + def INPUT_TYPES(cls) -> Dict[str, Any]: + """ + Defines the required input types. + + Returns: + Dict[str, Any]: An empty dictionary since this function requires no inputs. + """ + return { + "required": { + "encode": ("PTCALLABLE", {}), + "text_list": ("PYLIST", {}) + } + } + + RETURN_NAMES: Tuple[str, ...] = ("token_tens", "mask_tens") + RETURN_TYPES: Tuple[str, ...] = ("TENSOR", "TENSOR") + FUNCTION: str = "f" + CATEGORY: str = "Data Analysis" + + def f(self, + encode: Callable, + text_list: Union[str, List[str]]) -> Tuple[Callable]: + """ + Args: + encode (Callable): The reference to a token encoder function. + text_list (Union[str, List[str]]): The a string or a python list of strings to tokenize. + + Returns: + Tuple[Union[torch.Tensor, List[torch.Tensor]], Union[torch.Tensor, List[torch.Tensor]]]: A tuple containing the token IDs and the mask in torch.Tensor. + """ + return encode(text_list) diff --git a/modules/pytorch_wrapper/pt_train_rnn_model.py b/modules/pytorch_wrapper/pt_train_rnn_model.py index 558dff3..aca43e9 100644 --- a/modules/pytorch_wrapper/pt_train_rnn_model.py +++ b/modules/pytorch_wrapper/pt_train_rnn_model.py @@ -143,7 +143,10 @@ class PtTrainRNNModel: total_samples = len(train_loader.dataset) for x, y in train_loader: if use_gpu: - x = x.to("cuda") + if isinstance(x, list): + x = [e.to("cuda") for e in x] + else: + x = x.to("cuda") y = y.to("cuda") optimizer.zero_grad() if linear_head: @@ -192,7 +195,10 @@ class PtTrainRNNModel: total_samples = len(val_loader.dataset) for x, y in val_loader: if use_gpu: - x = x.to("cuda") + if isinstance(x, list): + x = [e.to("cuda") for e in x] + else: + x = x.to("cuda") y = y.to("cuda") if linear_head: y_hat = model(x) diff --git a/modules/pytorch_wrapper/ptn_embedding.py b/modules/pytorch_wrapper/ptn_embedding.py new file mode 100644 index 0000000..a6be914 --- /dev/null +++ b/modules/pytorch_wrapper/ptn_embedding.py @@ -0,0 +1,65 @@ +from typing import Any, Dict, Tuple +import torch +import torch.nn as nn + + +class PtnEmbedding: + """ + Constructs an embedding layer. + + Args: + num_embeddings (int): Vocabulary size. + embedding_dim (int): The size of each embedding vector. + padding_idx (int): Token ID to treat as padding. If a non-negative integer is specified, the embedding at that index will be excluded from gradient updates. Set to -1 to disable padding behavior (default). + + 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_embeddings": ("INT", {"default": 10000, "min": 1, "max": 1e6}), + "embedding_size": ("INT", {"default": 100, "min": 1, "max": 1e6}), + "padding_idx": ("INT", {"default": -1, "min": -1, "max": 1e6}), + } + } + + RETURN_TYPES: Tuple[str] = ("PTMODEL",) + FUNCTION: str = "f" + CATEGORY: str = "Training" + + @classmethod + def IS_CHANGED(cls, **kw): + return float("NaN") + + def f(self, + num_embeddings: int, + embedding_dim: int, + padding_idx: int + ) -> Tuple[nn.Module]: + """ + Constructs an embedding layer. + + Args: + num_embeddings (int): Vocabulary size. + embedding_dim (int): The size of each embedding vector. + padding_idx (int): Token ID to treat as padding. If a non-negative integer is specified, the embedding at that index will be excluded from gradient updates. Set to -1 to disable padding behavior (default). + + Returns: + Tuple[nn.Module]: A tuple containing the instantiated PyTorch model. + """ + + with torch.inference_mode(False): + model = nn.Embedding( + num_embeddings=num_embeddings, + embedding_dim=embedding_dim, + padding_idx=padding_idx) + + return (model,) diff --git a/modules/pytorch_wrapper/ptn_embedding_rnn_linear.py b/modules/pytorch_wrapper/ptn_embedding_rnn_linear.py new file mode 100644 index 0000000..6b45218 --- /dev/null +++ b/modules/pytorch_wrapper/ptn_embedding_rnn_linear.py @@ -0,0 +1,197 @@ +from typing import Any, Dict, Tuple +import torch +import torch.nn as nn +from pytorch_wrapper.rnn_output_reshape import extract_rnn_return_value_with_adjusted_label + + +class EmbeddingRNNLinear(nn.Module): + """ + A recurrent neural network (RNN) with a linear output layer. + + This module uses a standard PyTorch RNN, followed by a linear layer applied to the final output tensor. + + pragma: skip_doc + """ + + def __init__(self, + vocabulary_size: int, + input_size: int, + hidden_size: int, + num_layers: int, + nonlinearity: str, + bias: bool, + batch_first: bool, + dropout: float, + bidirectional: bool, + linear_output_size: int, + linear_bias: bool + ): + super().__init__() + self.embedding = nn.Embedding( + vocabulary_size, + input_size + ) + + self.rnn_model = nn.RNN( + input_size=input_size, + hidden_size=hidden_size, + num_layers=num_layers, + nonlinearity=nonlinearity, + bias=bias, + batch_first=batch_first, + dropout=dropout, + bidirectional=bidirectional + ) + self.bidirectional = bidirectional + self.batch_first = batch_first + directions = 2 if bidirectional else 1 + + # Initialize RNN weights + for layer in range(num_layers): + for direction in range(directions): + suffix = f"_reverse" if direction == 1 else "" + weight_ih = getattr(self.rnn_model, f"weight_ih_l{layer}{suffix}") + weight_hh = getattr(self.rnn_model, f"weight_hh_l{layer}{suffix}") + + if nonlinearity == "relu": + torch.nn.init.kaiming_normal_(weight_ih, nonlinearity='relu') + torch.nn.init.kaiming_normal_(weight_hh, nonlinearity='relu') + else: + torch.nn.init.xavier_normal_(weight_ih) + torch.nn.init.xavier_normal_(weight_hh) + + self.linear_model = nn.Linear( + in_features=hidden_size * directions, + out_features=linear_output_size, + bias=linear_bias + ) + torch.nn.init.xavier_uniform_(self.linear_model.weight) + + def forward(self, inputs: torch.Tensor) -> torch.Tensor: + """ + Forward pass through the RNN and the linear layer. + + Args: + inputs (torch.Tensor): Input tensor of shape (batch, seq, input_size) if batch_first=True. + + Returns: + torch.Tensor: Output of the linear layer. + """ + tokens, masks = inputs + embeds = self.embedding(tokens) + out_seq, _ = self.rnn_model(embeds) + + x = extract_rnn_return_value_with_adjusted_label( + masks.to(torch.float32).unsqueeze(-1), # This is needed to find out padded tokens + None, + out_seq, + self.bidirectional, + self.batch_first, + False, + return_valid_token_mean=True + ) + + outputs = self.linear_model(x) + return outputs + + +class PtnEmbeddingRNNLinear: + """ + Ptn RNN Linear: + A recurrent neural network (RNN) model with a linear head. + + Args: + vocabulary_size (int): The vocabulary size of the tokens. + input_size (int): The number of input features. + hidden_size (int): The number of output features of the hidden layer matrix. + num_layers (int): Number of hidden layers in RNN. + nonlinearity (str): Activation function to apply after each hidden layer. Specify 'tanh' or 'relu'. + bias (bool): Use bias or not. + batch_first (bool): Indicates the shape of the input. Input is assumed to be a rank-3 tensor. If True, the shape is [Batch, Seq, Token]. If False, the shape is [Seq, Batch, Token]. Note that default is set to `True` unlike PyTorch's RNN default model parameter. + dropout (float): Dropout probability to be applied to the output of intermediate layers if non-zero. This will not be applied to the last layer's output. + bidirectional (bool): Processes the input in a backward direction and append the result to the output of the forward direction. + linear_output_size (int): The number of output features of the linear layer. + linar_bias (bool): Use bias or not in 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": { + "vocabulary_size": ("INT", {"default": 10000, "min": 1, "max": 1e6}), + "input_size": ("INT", {"default": 1, "min": 1, "max": 1e6}), + "hidden_size": ("INT", {"default": 1, "min": 1, "max": 1e6}), + "num_layers": ("INT", {"default": 1, "min": 1, "max": 1e3}), + "nonlinearity": (("tanh", "relu"),), + "bias": ("BOOLEAN", {"default": True}), + "batch_first": ("BOOLEAN", {"default": True}), + "dropout": ("FLOAT", {"default": 0, "min": 0, "max": 1}), + "bidirectional": ("BOOLEAN", {"default": False}), + "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, + vocabulary_size: int, + input_size: int, + hidden_size: int, + num_layers: int, + nonlinearity: str, + bias: bool, + batch_first: bool, + dropout: float, + bidirectional: bool, + linear_output_size: int, + linear_bias: bool + ) -> Tuple[nn.Module]: + """ + Constructs an RNN model. + + Args: + vocabulary_size (int): The vocabulary size of the tokens. + input_size (int): The number of input features. + hidden_size (int): The number of output features of the hidden layer matrix. + num_layers (int): Number of hidden layers. + nonlinearity (str): Activation function to apply after each hidden layer. Specify 'tanh' or 'relu'. + bias (bool): Use bias or not. + batch_first (bool): Indicates the shape of the input. Input is assumed to be a rank-3 tensor. If True, the shape is [Batch, Seq, Token]. If False, the shape is [Seq, Batch, Token]. + dropout (float): Dropout probability to be applied to the output of intermediate layers if non-zero. This will not be applied to the last layer's output. + bidirectional (bool): Processes the input in a backward direction and append the result to the output of the forward direction. + linear_output_size (int): The number of output features of the linear layer. + linar_bias (bool): Use bias or not in linear layer. + + Returns: + Tuple[nn.Module]: A tuple containing the instantiated PyTorch model. + """ + with torch.inference_mode(False): + model = EmbeddingRNNLinear( + vocabulary_size=vocabulary_size, + input_size=input_size, + hidden_size=hidden_size, + num_layers=num_layers, + nonlinearity=nonlinearity, + bias=bias, + batch_first=batch_first, + dropout=dropout, + bidirectional=bidirectional, + linear_output_size=linear_output_size, + linear_bias=linear_bias + ) + return (model,) diff --git a/modules/pytorch_wrapper/ptv_hf_dataset_with_token_encode.py b/modules/pytorch_wrapper/ptv_hf_dataset_with_token_encode.py new file mode 100644 index 0000000..3183ba2 --- /dev/null +++ b/modules/pytorch_wrapper/ptv_hf_dataset_with_token_encode.py @@ -0,0 +1,75 @@ +from typing import Any, Dict, Tuple, Callable +import torch + + +class PtvHfDatasetWithTokenEncode: + """ + A PyTorch Dataset class wrapper for a Hugging Face dataset that converts text to token IDs. + + Args: + name (str): Name of the dataset class. + split (str): The name of the subset of dataset to retrieve or the split specifier. + 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. + + category: PyTorch wrapper - Training + """ + + @classmethod + def INPUT_TYPES(cls) -> Dict[str, Any]: + """ + Defines the required input types. + + Returns: + Dict[str, Any]: A dictionary of required input types. + """ + return { + "required": { + "name": ("STRING", {"default": "", "multiline": False}), + "split": ("STRING", {"default": "train","multiline": False}), + "sample_field_name": ("STRING", {"default": "text","multiline": False}), + "label_field_name": ("STRING", {"default": "label","multiline": False}), + "encode": ("PTCALLABLE", {}), + } + } + + RETURN_TYPES: Tuple[str] = ("PTVDATASET",) + FUNCTION: str = "f" + CATEGORY: str = "Training" + + @classmethod + def IS_CHANGED(cls, **kw): + return float("NaN") + + def f(self, + name: str, + split: str, + sample_field_name: str, + label_field_name: str, + encode: Callable) -> Tuple: + """ + Loads a dataset from Hugging Face with specified parameters. + + Args: + name (str): Name of the dataset class (e.g., MNIST, FashionMNIST). + split (str): The name of the subset of dataset to retrieve or the split specifier. + 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. + + 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, + split=split, + sample_field_name=sample_field_name, + label_field_name=label_field_name, + encode=encode) + return (dc,) diff --git a/modules/pytorch_wrapper/ptv_hf_glove_dataset.py b/modules/pytorch_wrapper/ptv_hf_glove_dataset.py index a37f1e8..6cd5660 100644 --- a/modules/pytorch_wrapper/ptv_hf_glove_dataset.py +++ b/modules/pytorch_wrapper/ptv_hf_glove_dataset.py @@ -7,7 +7,7 @@ class PtvHfGloveDataset: A PyTorch Dataset class wrapper for a Hugging Face dataset that converts text to embedding using Glove. Args: - name (str): Name of the dataset class (e.g., MNIST, FashionMNIST). + name (str): Name of the dataset class. split (str): The name of the subset of dataset to retrieve or the split specifier. embed_dim (int): Embedding dimension for Glove. max_seq_len (int): Sequence length. @@ -55,7 +55,7 @@ class PtvHfGloveDataset: Loads a dataset from Hugging Face with specified parameters. Args: - name (str): Name of the dataset class (e.g., MNIST, FashionMNIST). + name (str): Name of the dataset class. split (str): The name of the subset of dataset to retrieve or the split specifier. embed_dim (int): Embedding dimension for Glove. max_seq_len (int): Sequence length. diff --git a/modules/pytorch_wrapper/utils.py b/modules/pytorch_wrapper/utils.py index 9c5bf5d..333823a 100644 --- a/modules/pytorch_wrapper/utils.py +++ b/modules/pytorch_wrapper/utils.py @@ -1,6 +1,6 @@ import os import ast -from typing import Optional, Tuple, Union +from typing import Optional, Tuple, Union, List import random import numpy as np import torch @@ -75,7 +75,7 @@ def get_resource_full_path(resource_name: str, resource_path: str) -> str: return resource_full_path -def get_model_full_path(model_path: str) -> str: +def get_model_full_path(model_path: str, sub_model_dir_name: str=None) -> str: """ Resolves the absolute path to a model file within the "models" directory and ensures that the directory exists. @@ -83,11 +83,15 @@ def get_model_full_path(model_path: str) -> str: Args: model_path (str): The path to the model file. If the path is a relative path, the path is considered relative to the `models` directory of this extension. - + sub_model_dir_name (str): The subdirectory specific to the type of model (e.g. SentencePiece). Returns: str: The absolute path to the model file. """ - return get_resource_full_path("models", model_path) + if sub_model_dir_name: + model_dir = os.path.join("models", sub_model_dir_name) + else: + model_dir = "models" + return get_resource_full_path(model_dir, model_path) def get_dataset_full_path(dataset_path: str) -> str: @@ -152,4 +156,84 @@ def set_seed(seed=42) -> None: if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed) torch.backends.cudnn.deterministic = True - torch.backends.cudnn.benchmark = False \ No newline at end of file + torch.backends.cudnn.benchmark = False + + +def pad_truncate_sequence( + token_tensor_list: List[torch.Tensor], + padding: bool, + padding_method: str, + padding_value: Union[int, torch.Tensor], + truncation: bool, + max_length: int +) -> Union[ + Tuple[torch.Tensor, torch.Tensor], + Tuple[List[torch.Tensor], List[torch.Tensor]] +]: + """ + Pads and/or truncates a list of 1D token tensors to a uniform length. + + Args: + token_tensor_list (List[torch.Tensor]): List of 1D token tensors. + padding (bool): Whether to apply padding. + 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`. + padding_value (Union[int, torch.Tensor]): A scalar value or scalar tensor used for padding. + truncation (bool): Whether to truncate sequences to `max_length`. + max_length (int): The target length for truncation or padding. + + Returns: + Union[ + Tuple[torch.Tensor, torch.Tensor], # (tokens, masks) if padding is True or list has 1 element + Tuple[List[torch.Tensor], List[torch.Tensor]] # (tokens, masks) as lists otherwise + ] + """ + assert isinstance(token_tensor_list, list) and all(isinstance(t, torch.Tensor) for t in token_tensor_list), \ + "token_tensor_list must be a list of torch.Tensor" + if isinstance(padding_value, torch.Tensor): + assert padding_value.numel() == 1, "padding_value must be a scalar tensor" + padding_value = padding_value.item() + + if padding_method not in ("max_length", "longest"): + raise ValueError("Invalid padding_method. Use 'max_length' or 'longest'.") + + token_list = [] + mask_list = [] + + # Apply truncation and collect post-truncation lengths + truncated_list = [] + for t in token_tensor_list: + if t.dim() != 1: + raise ValueError("Expecting 1D tensor") + if truncation: + t = t[:max_length] # Out of bounds check is not needed as PyTorch handles that. + truncated_list.append(t) + + if padding and padding_method == "longest": + padding_target_length = max(t.size(0) for t in truncated_list) + elif padding and padding_method == "max_length": + padding_target_length = max_length + + for t in truncated_list: + tensor_length = t.size(0) + + # Apply padding if needed + if padding and tensor_length < padding_target_length: + filler_count = padding_target_length - tensor_length + pad_tensor = torch.full((filler_count,), padding_value, dtype=t.dtype, device=t.device) + t = torch.cat([t, pad_tensor], dim=0) + + # Create attention mask + if padding: + mask = (t != padding_value).long() + else: + mask = torch.ones_like(t) + + token_list.append(t) + mask_list.append(mask) + + if padding or len(token_tensor_list) == 1: + return torch.stack(token_list), torch.stack(mask_list) + else: + return token_list, mask_list diff --git a/modules/sentencepiece_wrapper/sp_encode.py b/modules/sentencepiece_wrapper/sp_encode.py new file mode 100644 index 0000000..bf0ec6c --- /dev/null +++ b/modules/sentencepiece_wrapper/sp_encode.py @@ -0,0 +1,134 @@ +import os +import sys +from typing import Any, Dict, Union, List, Tuple, Callable +import torch +import sentencepiece as spm + +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 pad_truncate_sequence + + +class SpEncode: + """ + SentencePiece wrapper for encoding text into token ID tensors. + + Args: + spmodel (SentencePieceProcessor): Trained SentencePiece 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`. + padding_value (Union[int, torch.Tensor]): A scalar value or scalar tensor used for padding. + truncation (bool): Whether to truncate sequences to `max_length`. + max_length (int): The target length for truncation or padding. + + 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": { + "spmodel": ("SPMODEL", {}), + "padding": ("BOOLEAN", {"default": False}), + "padding_method": (("max_length", "longest"),), + "padding_value": ("INT", {"default": 0, "min":0, "max":255}), + "truncation": ("BOOLEAN", {"default": False}), + "max_length": ("INT", {"default": 512, "min": 1, "max": 1_000_000}), + } + } + + RETURN_TYPES: Tuple[str] = ("PTCALLABLE",) + FUNCTION: str = "f" + CATEGORY: str = "Data Analysis" + + def encode( + self, + sentence: Union[str, List[str]], + ) -> Tuple[Union[torch.Tensor, List[torch.Tensor]], Union[torch.Tensor, List[torch.Tensor]]]: + """ + Encodes a string or list of strings into token ID tensors. + + Args: + sentence (Union[str, List[str]]): Input sentence(s). + + Returns: + Tuple of token tensors and attention masks: + - If padding is True: + (2D token tensor, 2D attention mask tensor) + - If padding is False: + - For a single string: (1D token tensor, 1D mask tensor) + - For a list: (List of 1D token tensors, List of 1D mask tensors) + """ + if isinstance(sentence, str): + rank1_input = True + sentence = [sentence] + else: + rank1_input = False + + token_list = [] + for s in sentence: + ids = self.spmodel.encode(s, out_type=int) + if self.truncation and len(ids) > self.max_length: + ids = ids[:self.max_length] + token_list.append(torch.tensor(ids, dtype=torch.long)) + + tokens, masks = pad_truncate_sequence( + token_list, + self.padding, + self.padding_method, + self.padding_value, + self.truncation, + self.max_length + ) + + if rank1_input: + return tokens[0], masks[0] + else: + return tokens, masks + + def f( + self, + spmodel: spm.SentencePieceProcessor, + padding: bool, + padding_method: str, + padding_value: int, + truncation: bool, + max_length: int + ) -> Tuple[Callable[[Union[str, List[str]]], Tuple[Union[torch.Tensor, List[torch.Tensor]], Union[torch.Tensor, List[torch.Tensor]]]]]: + """ + Binds configuration and returns a callable encoder function. + + Args: + spmodel (SentencePieceProcessor): Trained SentencePiece 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`. + padding_value (Union[int, torch.Tensor]): A scalar value or scalar tensor used for padding. + truncation (bool): Whether to truncate sequences to `max_length`. + max_length (int): The target length for truncation or padding. + + Returns: + Tuple[Callable]: A tuple containing a callable that accepts a sentence or list of + sentences and returns token ID tensor(s) and attention mask(s). + """ + self.spmodel = spmodel + self.padding = padding + self.padding_method = padding_method + self.padding_value = padding_value + self.truncation = truncation + self.max_length = max_length + return (self.encode,) diff --git a/modules/sentencepiece_wrapper/sp_load_model.py b/modules/sentencepiece_wrapper/sp_load_model.py new file mode 100644 index 0000000..7259a42 --- /dev/null +++ b/modules/sentencepiece_wrapper/sp_load_model.py @@ -0,0 +1,47 @@ +import os +import sys +from typing import Any, Dict +import torch +import sentencepiece as spm + +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 get_model_full_path + +class SpLoadModel: + """ + A wrapper class for loading a SentencePiece tokenization model. + + 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_path": ("STRING", {"default": "spiece.model"}) + } + } + + RETURN_NAMES: tuple = ("spmodel",) + RETURN_TYPES: tuple = ("SPMODEL",) + FUNCTION: str = "f" + CATEGORY: str = "Data Analysis" + + @classmethod + def IS_CHANGED(cls, **kw): + return float("NaN") + + def f(self, model_path: str): + model_full_path = get_model_full_path(model_path, + sub_model_dir_name="sentencepiece") + sp = spm.SentencePieceProcessor(model_file=model_full_path) + + return (sp,) \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index 4cb6c8b..ed9d900 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.2" +version = "3.3.3" license = {file = "LICENSE"} [project.urls] diff --git a/requirements.txt b/requirements.txt index f308048..c05d0da 100644 --- a/requirements.txt +++ b/requirements.txt @@ -5,3 +5,4 @@ scipy>=1.14.1 scikit-learn>=1.6.1 datasets>=3.4.1 gensim==4.3.2 +sentencepiece>=0.2.0 diff --git a/tests/pytorch_wrapper/hf_dataset_with_tokenize_test.py b/tests/pytorch_wrapper/hf_dataset_with_tokenize_test.py new file mode 100644 index 0000000..51d9277 --- /dev/null +++ b/tests/pytorch_wrapper/hf_dataset_with_tokenize_test.py @@ -0,0 +1,62 @@ +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_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, + ) + + 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.") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/pytorch_wrapper/pt_tokenizer_test.py b/tests/pytorch_wrapper/pt_tokenizer_test.py new file mode 100644 index 0000000..fbb3c37 --- /dev/null +++ b/tests/pytorch_wrapper/pt_tokenizer_test.py @@ -0,0 +1,43 @@ +import os +import sys +import unittest +import sentencepiece as spm +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 sentencepiece_wrapper.sp_load_model import SpLoadModel +from sentencepiece_wrapper.sp_encode import SpEncode +from pytorch_wrapper.pt_tokenizer import PtTokenizer + +class TestPtTokenizer(unittest.TestCase): + def setUp(self): + self.load_node = SpLoadModel() + self.encode_node = SpEncode() + self.tokenizer_node = PtTokenizer() + + self.sp = self.load_node.f("spiece.model")[0] + + def test_1(self): + sentence = "Hello, world" + encode = self.encode_node.f( + spmodel=self.sp, + padding=False, + padding_method="longest", + padding_value=0, + truncation=False, + max_length=256 + )[0] + + tokens, mask = self.tokenizer_node.f(encode=encode, text_list=sentence) + + expected_tokens = torch.tensor([8774, 6, 296], dtype=torch.long) + expected_mask = torch.tensor([1, 1, 1], dtype=torch.long) + self.assertTrue(torch.equal(tokens, expected_tokens)) + self.assertTrue(torch.equal(mask, expected_mask)) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/pytorch_wrapper/pt_train_embedding_rnn_linear_text_classification_model_test.py b/tests/pytorch_wrapper/pt_train_embedding_rnn_linear_text_classification_model_test.py new file mode 100644 index 0000000..184dba8 --- /dev/null +++ b/tests/pytorch_wrapper/pt_train_embedding_rnn_linear_text_classification_model_test.py @@ -0,0 +1,133 @@ +""" +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_rnn_linear import PtnEmbeddingRNNLinear +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_adam import PtoAdam +from pytorch_wrapper.ptv_hf_glove_dataset import PtvHfGloveDataset +from pytorch_wrapper.pt_data_loader import PtDataLoader +from pytorch_wrapper.pt_train_rnn_model import PtTrainRNNModel +from pytorch_wrapper.pt_save_model import PtSaveModel +from pytorch_wrapper.pt_load_model import PtLoadModel +from pytorch_wrapper.pto_lr_scheduler_reduce_on_plateau import PtoLrSchedulerReduceOnPlateau +from pytorch_wrapper.ptn_bce_with_logits_loss import PtnBCEWithLogitsLoss +from pytorch_wrapper.utils import set_seed + + +class TestPtTrainEmbeddingRNNLinearTextClassificationModel(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=300 + 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, + ) + + test_dataset = HfDatasetWithTokenEncode( + "imdb", # dataset_name + "test", #split + encode=encode, + ) + + train_data_loader_node = PtDataLoader() + test_data_loader_node = PtDataLoader() + + self.train_loader = train_data_loader_node.f( + dataset=train_dataset, + batch_size=128, + shuffle=True, + parameters='{"num_workers":1}' + )[0] + + self.test_loader = test_data_loader_node.f( + dataset=test_dataset, + batch_size=128, + shuffle=False, + parameters='{"num_workers":1}' + )[0] + + self.model_path = "rnn_text_classification_test.pt" + + self.model_node = PtnEmbeddingRNNLinear() + self.model = self.model_node.f( + 32000, # vocabulary + 300, # input size + 128, # hidden_size + 2, # num_layers + "tanh", # nonlinearity + True, # bias, + True, # batch_first [B, Seq, Token] + 0.1, # dropout + True, # bidirectional, + 1, # linear output size + True # linear bias + )[0] + + self.optimizer = PtoAdam().f(self.model, 0.001, 0.9, 0.999)[0] + + self.scheduler = PtoLrSchedulerReduceOnPlateau().f(self.optimizer, + grace_period=3, + gamma=0.5 + )[0] + + self.trainer = PtTrainRNNModel() + self.save_model = PtSaveModel() + self.load_model = PtLoadModel() + self.linear_head = True + self.epochs = 50 + + @unittest.skipIf(os.getenv("RUN_SKIPPED_TESTS") != "1", "Skipping unless explicitly enabled") + def test_1(self): + + trained_model = self.trainer.f( + self.model, + self.linear_head, + 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, + use_valid_token_mean=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_rnn_linear_test.py b/tests/pytorch_wrapper/ptn_embedding_rnn_linear_test.py new file mode 100644 index 0000000..18fd6f8 --- /dev/null +++ b/tests/pytorch_wrapper/ptn_embedding_rnn_linear_test.py @@ -0,0 +1,41 @@ +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_rnn_linear import PtnEmbeddingRNNLinear + + +class TestPtnEmbeddingRNNLinear(unittest.TestCase): + def setUp(self): + """Set up test instance.""" + self.model_node = PtnEmbeddingRNNLinear() + + def test_1(self): + self.model = self.model_node.f( + 32000, # vocabulary_size + 2, # input_size + 3, # hidden_size + 1, # num_layers + "tanh", # nonlinearity + True, # bias, + True, # batch_first + 0.0, # dropout + False, # bidirectional, + 5, # linear output features + True # linear bias + )[0] + + x = torch.randint(0, 32000, (8 * 5,)).to(torch.int64) + x = x.view(8, 5) + masks = torch.ones_like(x) + outputs = self.model((x, masks)) + self.assertEqual(outputs.size(), torch.Size([8, 5])) + + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/pytorch_wrapper/ptn_embedding_test.py b/tests/pytorch_wrapper/ptn_embedding_test.py new file mode 100644 index 0000000..5c99ea5 --- /dev/null +++ b/tests/pytorch_wrapper/ptn_embedding_test.py @@ -0,0 +1,61 @@ +import os +import sys +import unittest +import torch +import torch.nn as nn +from typing import Any + +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 import PtnEmbedding + + +class TestPtnEmbedding(unittest.TestCase): + def setUp(self) -> None: + """Set up the test instance of PtnEmbedding.""" + self.model_node = PtnEmbedding() + + def test_basic_embedding(self) -> None: + """Test that PtnEmbedding outputs expected shapes and has correct weight dimensions.""" + # Create embedding module + model: nn.Module = self.model_node.f( + 100, # vocabulary size + 8, # embedding dimension + -1 # no padding index + )[0] + + # Check embedding weight shape + self.assertEqual(model.weight.size(), torch.Size([100, 8])) + + # Create sample input: batch size 1, sequence length 6 + x: torch.Tensor = torch.tensor([[0, 1, 2, 3, 4, 5]], dtype=torch.int64) + + # Forward pass + out: torch.Tensor = model(x) + + # Output shape should be (batch_size, sequence_length, embedding_dim) + self.assertEqual(out.size(), torch.Size([1, 6, 8])) + + + def test_padding_index(self) -> None: + """Test that PtnEmbedding sets the correct padding_idx internally.""" + padding_idx: int = 2 + model: nn.Module = self.model_node.f( + 100, # vocabulary size + 8, # embedding dimension + padding_idx + )[0] + + # Ensure the model has the padding_idx attribute set correctly + self.assertTrue(hasattr(model, "padding_idx")) + self.assertEqual(model.padding_idx, padding_idx) + + # Check that the weights at padding_idx are all zeros initially + self.assertTrue(torch.all(model.weight[padding_idx] == 0)) + + + +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 new file mode 100644 index 0000000..f674719 --- /dev/null +++ b/tests/pytorch_wrapper/ptv_hf_dataset_with_token_encode_test.py @@ -0,0 +1,65 @@ +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.ptv_hf_dataset_with_token_encode import PtvHfDatasetWithTokenEncode +from sentencepiece_wrapper.sp_load_model import SpLoadModel +from sentencepiece_wrapper.sp_encode import SpEncode + +class TestPtvHfDatasetWithTokenEncode(unittest.TestCase): + """ + Tests the node for Hugging Face Dataset wrapper with word embedding support. + """ + + 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 = PtvHfDatasetWithTokenEncode().f( + "imdb", # dataset_name + "train", #split + "text", + "label", + encode + )[0] + + + 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.") + + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/pytorch_wrapper/utils_pad_truncate_sequence_test.py b/tests/pytorch_wrapper/utils_pad_truncate_sequence_test.py new file mode 100644 index 0000000..f23d6b4 --- /dev/null +++ b/tests/pytorch_wrapper/utils_pad_truncate_sequence_test.py @@ -0,0 +1,120 @@ +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 pytorch_wrapper.utils import pad_truncate_sequence + + +class TestPadTruncateSequence(unittest.TestCase): + def setUp(self): + # 3 sample sequences of different lengths + self.token_tensor_list = [ + torch.tensor([1, 2, 3], dtype=torch.long), + torch.tensor([4, 5], dtype=torch.long), + torch.tensor([6], dtype=torch.long) + ] + self.max_length = 255 + + def run_case(self, padding, padding_method, padding_value, truncation): + + with self.subTest( + padding=padding, + padding_method=padding_method, + padding_value=padding_value, + truncation=truncation + ): + tokens, masks = pad_truncate_sequence( + token_tensor_list=self.token_tensor_list, + padding=padding, + padding_method=padding_method, + padding_value=padding_value, + truncation=truncation, + max_length=self.max_length + ) + + if padding: + # Tokens and masks should be stacked tensors + self.assertIsInstance(tokens, torch.Tensor) + self.assertIsInstance(masks, torch.Tensor) + self.assertEqual(tokens.shape, masks.shape) + self.assertEqual(tokens.dim(), 2) # batch x seq_len + self.assertEqual(tokens.size(0), 3) + if padding_method == "max_length": + self.assertEqual(tokens.size(1), self.max_length) + elif padding_method == "longest": + truncated_lengths = [min(len(t), self.max_length) if truncation else len(t) + for t in self.token_tensor_list] + expected_len = max(truncated_lengths) + self.assertEqual(tokens.size(1), expected_len) + else: + # Should return lists + self.assertIsInstance(tokens, list) + self.assertIsInstance(masks, list) + self.assertEqual(len(tokens), 3) + for t, m in zip(tokens, masks): + self.assertEqual(t.size(), m.size()) + expected_len = min(len(t), self.max_length) if truncation else len(t) + self.assertEqual(t.size(0), expected_len) + + def test_all_cases(self): + for padding, padding_method, padding_value, truncation in product( + [True, False], + ["max_length", "longest"], + [0, 255], + [True, False] + ): + self.run_case(padding, padding_method, padding_value, truncation) + + def test_single_tensor_input(self): + """Test behavior when token_tensor_list has only one element.""" + single_tensor_list = [torch.tensor([1, 2, 3, 4], dtype=torch.long)] + + for padding, padding_method, padding_value, truncation in product( + [True, False], + ["max_length", "longest"], + [0, 255], + [True, False] + ): + with self.subTest( + single_input=True, + padding=padding, + padding_method=padding_method, + padding_value=padding_value, + truncation=truncation + ): + tokens, masks = pad_truncate_sequence( + token_tensor_list=single_tensor_list, + padding=padding, + padding_method=padding_method, + padding_value=padding_value, + truncation=truncation, + max_length=self.max_length + ) + + # Always stacked even if padding is False + self.assertIsInstance(tokens, torch.Tensor) + self.assertIsInstance(masks, torch.Tensor) + self.assertEqual(tokens.shape, masks.shape) + self.assertEqual(tokens.size(0), 1) + + expected_len = 4 + if truncation: + expected_len = min(expected_len, self.max_length) + if padding: + if padding_method == "max_length": + expected_len = self.max_length + elif padding_method == "longest": + expected_len = expected_len # no change + + self.assertEqual(tokens.size(1), expected_len) + + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/sentencepiece_wrapper/sp_encode_test.py b/tests/sentencepiece_wrapper/sp_encode_test.py new file mode 100644 index 0000000..8ef6b11 --- /dev/null +++ b/tests/sentencepiece_wrapper/sp_encode_test.py @@ -0,0 +1,146 @@ +import os +import sys +import unittest +import sentencepiece as spm +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 sentencepiece_wrapper.sp_load_model import SpLoadModel +from sentencepiece_wrapper.sp_encode import SpEncode + + +class TestSpEncode(unittest.TestCase): + def setUp(self): + self.load_node = SpLoadModel() + self.node = SpEncode() + self.sp = self.load_node.f("spiece.model")[0] + + def test_1d_no_padding(self): + sentence = "Hello, world" + encoder = self.node.f( + spmodel=self.sp, + padding=False, + padding_method="longest", + padding_value=0, + truncation=False, + max_length=256 + )[0] + tokens, mask = encoder(sentence=sentence) + print(tokens) + expected_tokens = torch.tensor([8774, 6, 296], dtype=torch.long) + expected_mask = torch.tensor([1, 1, 1], dtype=torch.long) + self.assertTrue(torch.equal(tokens, expected_tokens)) + self.assertTrue(torch.equal(mask, expected_mask)) + + def test_2d_no_padding(self): + sentence = ["Hello,", "world"] + encoder = self.node.f( + spmodel=self.sp, + padding=False, + padding_method="longest", + padding_value=0, + truncation=False, + max_length=256 + )[0] + tokens, masks = encoder(sentence=sentence) + + expected_tokens = [ + torch.tensor([8774, 6], dtype=torch.long), + torch.tensor([296], dtype=torch.long) + ] + expected_masks = [ + torch.tensor([1, 1], dtype=torch.long), + torch.tensor([1], dtype=torch.long) + ] + for t, e in zip(tokens, expected_tokens): + self.assertTrue(torch.equal(t, e)) + for m, e in zip(masks, expected_masks): + self.assertTrue(torch.equal(m, e)) + + def test_1d_with_padding(self): + sentence = "Hello, world" + encoder = self.node.f( + spmodel=self.sp, + padding=True, + padding_method="longest", + padding_value=0, + truncation=False, + max_length=256 + )[0] + tokens, mask = encoder(sentence=sentence) + expected_tokens = torch.tensor([8774, 6, 296], dtype=torch.long) + expected_mask = torch.tensor([1, 1, 1], dtype=torch.long) + self.assertTrue(torch.equal(tokens, expected_tokens)) + self.assertTrue(torch.equal(mask, expected_mask)) + + def test_2d_with_padding(self): + sentence = ["Hello,", "world"] + encoder = self.node.f( + spmodel=self.sp, + padding=True, + padding_method="longest", + padding_value=0, + truncation=False, + max_length=256 + )[0] + tokens, mask = encoder(sentence=sentence) + + expected_tokens = torch.tensor([ + [8774, 6], + [296, 0] + ], dtype=torch.long) + expected_mask = torch.tensor([ + [1, 1], + [1, 0] + ], dtype=torch.long) + self.assertTrue(torch.equal(tokens, expected_tokens)) + self.assertTrue(torch.equal(mask, expected_mask)) + + def test_1d_truncation(self): + sentence = "Hello, world" + encoder = self.node.f( + spmodel=self.sp, + padding=False, + padding_method="longest", + padding_value=0, + truncation=True, + max_length=2 + )[0] + tokens, mask = encoder(sentence=sentence) + + expected_tokens = torch.tensor([8774, 6], dtype=torch.long) + expected_mask = torch.tensor([1, 1], dtype=torch.long) + self.assertTrue(torch.equal(tokens, expected_tokens)) + self.assertTrue(torch.equal(mask, expected_mask)) + + def test_2d_truncation(self): + sentence = ["Hello,", "world"] + encoder = self.node.f( + spmodel=self.sp, + padding=False, + padding_method="longest", + padding_value=0, + truncation=True, + max_length=1 + )[0] + tokens, masks = encoder(sentence=sentence) + + expected_tokens = [ + torch.tensor([8774], dtype=torch.long), + torch.tensor([296], dtype=torch.long) + ] + expected_masks = [ + torch.tensor([1], dtype=torch.long), + torch.tensor([1], dtype=torch.long) + ] + for t, e in zip(tokens, expected_tokens): + self.assertTrue(torch.equal(t, e)) + for m, e in zip(masks, expected_masks): + self.assertTrue(torch.equal(m, e)) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/sentencepiece_wrapper/sp_load_model_test.py b/tests/sentencepiece_wrapper/sp_load_model_test.py new file mode 100644 index 0000000..ab20ea0 --- /dev/null +++ b/tests/sentencepiece_wrapper/sp_load_model_test.py @@ -0,0 +1,29 @@ +import os +import sys +import unittest +import sentencepiece as spm + +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 sentencepiece_wrapper.sp_load_model import SpLoadModel + + +class TestSpLoadModel(unittest.TestCase): + def setUp(self): + """Set up test data.""" + self.node = SpLoadModel() + + + def test_1(self): + """Tests loading.""" + sp = self.node.f("spiece.model")[0] + + # Check tokenization + s = "Hello, world" + actual = sp.encode(s, out_type=int) # Token IDs + expected = [8774, 6, 296] + self.assertEqual(expected, actual, f"expected {expected} and actual {actual} do not match.") + +if __name__ == "__main__": + unittest.main()