nodes/HowToSD-ComfyUI-Pt-Wrapper
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Added Embedding layer and Embedding RNN Layer node.

Browse Source
This commit is contained in:
HowToSD
2025-03-27 11:51:21 -07:00
parent 8488471e8f
commit faf3a8f0e5
33 changed files with 1899 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files
__init__.py
+18
View File
@@ -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",
}
credit/credit.md
+254
View File
@@ -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 <a href="https://apache.org/foundation/license-faq.html#Apply-My-Software">recommends</a> 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.
```
docs/reference/node_reference.md
+6
View File
@@ -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 |
| --- | --- |
docs/reference/pt_tokenizer.md
+26
View File
@@ -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 |
<HR>
Category: PyTorch wrapper - Callable
ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
docs/reference/ptn_embedding.md
+24
View File
@@ -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 |
<HR>
Category: PyTorch wrapper - Model
ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
docs/reference/ptn_embedding_rnn_linear.md
+40
View File
@@ -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 |
<HR>
Category: PyTorch wrapper - Model
ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
docs/reference/ptv_hf_dataset_with_token_encode.md
+28
View File
@@ -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 |
<HR>
Category: PyTorch wrapper - Training
ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
docs/reference/ptv_hf_glove_dataset.md
+1 -1
View File
@@ -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.
docs/reference/sp_encode.md
+36
View File
@@ -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 |
<HR>
Category: PyTorch wrapper - Training
ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
docs/reference/sp_load_model.md
+17
View File
@@ -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 |
<HR>
Category: PyTorch wrapper - Training
ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
docs/training_rnn_for_classification.md
+18 -1
View File
@@ -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.
examples/workflows/embedding_rnn_classification.json
+1
View File
File diff suppressed because one or more lines are too long
modules/pytorch_wrapper/hf_dataset_with_token_encode.py
+62
View File
@@ -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)
modules/pytorch_wrapper/pt_tokenizer.py
+49
View File
@@ -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)
modules/pytorch_wrapper/pt_train_rnn_model.py
+8 -2
View File
@@ -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)
modules/pytorch_wrapper/ptn_embedding.py
+65
View File
@@ -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,)
modules/pytorch_wrapper/ptn_embedding_rnn_linear.py
+197
View File
@@ -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,)
modules/pytorch_wrapper/ptv_hf_dataset_with_token_encode.py
+75
View File
@@ -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,)
modules/pytorch_wrapper/ptv_hf_glove_dataset.py
+2 -2
View File
@@ -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.
modules/pytorch_wrapper/utils.py
+89 -5
View File
@@ -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
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
modules/sentencepiece_wrapper/sp_encode.py
+134
View File
@@ -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,)
modules/sentencepiece_wrapper/sp_load_model.py
+47
View File
@@ -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,)
pyproject.toml
+1 -1
View File
@@ -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]
requirements.txt
+1
View File
@@ -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
tests/pytorch_wrapper/hf_dataset_with_tokenize_test.py
+62
View File
@@ -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()
tests/pytorch_wrapper/pt_tokenizer_test.py
+43
View File
@@ -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()
tests/pytorch_wrapper/pt_train_embedding_rnn_linear_text_classification_model_test.py
+133
View File
@@ -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()
tests/pytorch_wrapper/ptn_embedding_rnn_linear_test.py
+41
View File
@@ -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()
tests/pytorch_wrapper/ptn_embedding_test.py
+61
View File
@@ -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()
tests/pytorch_wrapper/ptv_hf_dataset_with_token_encode_test.py
+65
View File
@@ -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()
tests/pytorch_wrapper/utils_pad_truncate_sequence_test.py
+120
View File
@@ -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()
tests/sentencepiece_wrapper/sp_encode_test.py
+146
View File
@@ -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()
tests/sentencepiece_wrapper/sp_load_model_test.py
+29
View File
@@ -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()