From 81bed6f3c5c419d043ba4d863d5909e001831b8a Mon Sep 17 00:00:00 2001 From: HowToSD Date: Tue, 1 Apr 2025 17:10:47 -0700 Subject: [PATCH] Added a node to load a dataset from the local file system. --- __init__.py | 3 + dev/scripts/save_hf_dataset_locally.py | 35 +++++++ ..._hugging_face_models_for_classification.md | 9 +- docs/reference/node_reference.md | 1 + docs/reference/ptv_hf_local_dataset.md | 33 +++++++ ...stilbert_classification_local_dataset.json | 1 + .../hf_dataset_with_token_encode.py | 13 +-- .../hugging_face_wrapper/hf_local_dataset.py | 95 +++++++++++++++++++ modules/hugging_face_wrapper/utils.py | 27 ++++++ .../pytorch_wrapper/ptv_hf_local_dataset.py | 90 ++++++++++++++++++ pyproject.toml | 2 +- ...h_tokenize_with_hf_tokenize_encode_test.py | 63 ++++++++++++ tests/hugging_face_wrapper/utils_test.py | 29 ++++++ .../ptv_hf_local_dataset_test.py | 68 +++++++++++++ 14 files changed, 456 insertions(+), 13 deletions(-) create mode 100644 dev/scripts/save_hf_dataset_locally.py create mode 100644 docs/reference/ptv_hf_local_dataset.md create mode 100644 examples/workflows/distilbert_classification_local_dataset.json create mode 100644 modules/hugging_face_wrapper/hf_local_dataset.py create mode 100644 modules/hugging_face_wrapper/utils.py create mode 100644 modules/pytorch_wrapper/ptv_hf_local_dataset.py create mode 100644 tests/hugging_face_wrapper/hf_local_dataset_with_tokenize_with_hf_tokenize_encode_test.py create mode 100644 tests/hugging_face_wrapper/utils_test.py create mode 100644 tests/pytorch_wrapper/ptv_hf_local_dataset_test.py diff --git a/__init__.py b/__init__.py index 717a5fc..0c06e6c 100644 --- a/__init__.py +++ b/__init__.py @@ -197,6 +197,7 @@ 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_hf_local_dataset import PtvHfLocalDataset 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 @@ -410,6 +411,7 @@ NODE_CLASS_MAPPINGS: Dict[str, Type[T]] = { "PtvDatasetLoader": PtvDatasetLoader, "PtvHfDatasetWithTokenEncode": PtvHfDatasetWithTokenEncode, "PtvHfGloveDataset": PtvHfGloveDataset, + "PtvHfLocalDataset": PtvHfLocalDataset, "PtvImageFolderDataset": PtvImageFolderDataset, "PtvSequentialTensorDataset": PtvSequentialTensorDataset, "PtvTransformsDataAugment": PtvTransformsDataAugment, @@ -623,6 +625,7 @@ NODE_DISPLAY_NAME_MAPPINGS: Dict[str, str] = { "PtvDatasetLoader": "Ptv Dataset Loader", "PtvHfDatasetWithTokenEncode": "Ptv Hf Dataset With Token Encode", "PtvHfGloveDataset": "Ptv Hf Glove Dataset", + "PtvHfLocalDataset": "Ptv Hf Local Dataset", "PtvImageFolderDataset": "Ptv Image Folder Dataset", "PtvSequentialTensorDataset": "Ptv Sequential Tensor Dataset", "PtvTransformsDataAugment": "Ptv Transforms Data Augment", diff --git a/dev/scripts/save_hf_dataset_locally.py b/dev/scripts/save_hf_dataset_locally.py new file mode 100644 index 0000000..171f33d --- /dev/null +++ b/dev/scripts/save_hf_dataset_locally.py @@ -0,0 +1,35 @@ +import os +import json +from datasets import load_dataset + +SOURCE_DATASET="tbd" # Update this line to match the actual dataset. +TARGET_DIR = os.path.join( + os.path.dirname(__file__), + "..", "..", # project root + "datasets", SOURCE_DATASET) +TRAIN_PATH = os.path.join(TARGET_DIR, "train.jsonl") +TEST_PATH = os.path.join(TARGET_DIR, "test.jsonl") + +def main(): + + # Check target directory + if os.path.exists(TARGET_DIR) is False: + os.makedirs(TARGET_DIR) + print(f"Created {TARGET_DIR}") + + # Load the dataset + dataset = load_dataset(SOURCE_DATASET) + + # Save train split + with open(TRAIN_PATH, "w") as f: + for item in dataset["train"]: + f.write(json.dumps(item) + "\n") + + # Save test split + with open(TEST_PATH, "w") as f: + for item in dataset["test"]: + f.write(json.dumps(item) + "\n") + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/docs/fine_tuning_hugging_face_models_for_classification.md b/docs/fine_tuning_hugging_face_models_for_classification.md index d6a7deb..ccabdd3 100644 --- a/docs/fine_tuning_hugging_face_models_for_classification.md +++ b/docs/fine_tuning_hugging_face_models_for_classification.md @@ -68,4 +68,11 @@ Also specify the total number of epochs in `epochs`. For example, if you set `freeze_pretrained_module_epochs` to 1 and `epochs` to 4, then the weights for the pretrained module are frozen during the first epoch, and full fine-tuning takes place from epochs 2 to 4. -Using this example workflow, you can expect a validation accuracy of 92% or 93%. \ No newline at end of file +Using this example workflow, you can expect a validation accuracy of 92% or 93%. + +# Loading a Dataset from the Local File System + +Instead of downloading a dataset from Hugging Face, you can also specify a custom dataset stored on your local file system. +Use the **PtvHfLocalDataset** node to specify the file path for the dataset. + +See `examples/workflows/distilbert_classification_local_dataset.json` for an example workflow that demonstrates loading an IMDB dataset in JSONL format from the local file system. diff --git a/docs/reference/node_reference.md b/docs/reference/node_reference.md index 5bba24c..537f308 100644 --- a/docs/reference/node_reference.md +++ b/docs/reference/node_reference.md @@ -251,6 +251,7 @@ | [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 Hf Local Dataset](ptv_hf_local_dataset.md) | A PyTorch Dataset class wrapper for a Hugging Face dataset to load a dataset that is stored on the local file system. | | [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. | diff --git a/docs/reference/ptv_hf_local_dataset.md b/docs/reference/ptv_hf_local_dataset.md new file mode 100644 index 0000000..e5331c5 --- /dev/null +++ b/docs/reference/ptv_hf_local_dataset.md @@ -0,0 +1,33 @@ +# Ptv Hf Local Dataset +A PyTorch Dataset class wrapper for a Hugging Face dataset to load a dataset that is stored on the local file system. + + Args: + file_path (str): Path to the dataset file. Can be an absolute path or a relative path. + If a relative path is provided, it is resolved from the `datasets` directory under this extension. + For example, to load `datasets/foo/bar_train.jsonp`, specify `foo/bar_train.jsonp` or the absolute path. + file_format (str): File format (e.g. json if the file is in jsonl format). sample_field_name (str): Field name for text samples. + label_field_name (str): Field name for labels. + encode (PTCALLABLE): The reference to a token encoder function. + remove_html_tags (bool): Remove html tags in text if True. + encode_return_dict (bool): Encode function returns a Dict instead of a Tuple. + +## Input +| Name | Data type | +|---|---| +| file_path | String | +| file_format | String | +| sample_field_name | String | +| label_field_name | String | +| encode | Ptcallable | +| remove_html_tags | Boolean | +| encode_return_dict | Boolean | + +## Output +| Data type | +|---| +| Ptvdataset | + +
+Category: PyTorch wrapper - Training + +ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved. diff --git a/examples/workflows/distilbert_classification_local_dataset.json b/examples/workflows/distilbert_classification_local_dataset.json new file mode 100644 index 0000000..8b05efc --- /dev/null +++ b/examples/workflows/distilbert_classification_local_dataset.json @@ -0,0 +1 @@ +{"last_node_id":37,"last_link_id":53,"nodes":[{"id":19,"type":"PandasShowText","pos":[3504.62841796875,1897.03662109375],"size":[302.7994079589844,112],"flags":{},"order":13,"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/distilbert_classification_local_dataset.pt"]},{"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":[36],"slot_index":0}],"properties":{"Node name for S&R":"PtnBCEWithLogitsLoss"},"widgets_values":["mean"]},{"id":7,"type":"PtDataLoader","pos":[2293.7060546875,1735.56494140625],"size":[271.68658447265625,173.3039093017578],"flags":{},"order":9,"mode":0,"inputs":[{"name":"dataset","localized_name":"dataset","type":"PTVDATASET","link":52}],"outputs":[{"name":"PTDATALOADER","localized_name":"PTDATALOADER","type":"PTDATALOADER","links":[34],"slot_index":0}],"properties":{"Node name for S&R":"PtDataLoader"},"widgets_values":[32,true,"{\"num_workers\":1}"]},{"id":12,"type":"PtDataLoader","pos":[2290.658447265625,2016.2130126953125],"size":[271.68658447265625,173.3039093017578],"flags":{},"order":10,"mode":0,"inputs":[{"name":"dataset","localized_name":"dataset","type":"PTVDATASET","link":53}],"outputs":[{"name":"PTDATALOADER","localized_name":"PTDATALOADER","type":"PTDATALOADER","links":[35],"slot_index":0}],"properties":{"Node name for S&R":"PtDataLoader"},"widgets_values":[32,false,"{\"num_workers\":1}"]},{"id":32,"type":"PtnHfFineTunedClassificationModel","pos":[2212.781494140625,1116.2008056640625],"size":[319.20001220703125,109.10855102539062],"flags":{},"order":5,"mode":0,"inputs":[{"name":"model_name","type":"STRING","widget":{"name":"model_name"},"link":48}],"outputs":[{"name":"PTMODEL","localized_name":"PTMODEL","type":"PTMODEL","links":[47,49],"slot_index":0}],"properties":{"Node name for S&R":"PtnHfFineTunedClassificationModel"},"widgets_values":["",true,0]},{"id":29,"type":"PtoAdamW","pos":[2673.749267578125,1107.9647216796875],"size":[285.66900634765625,158.52035522460938],"flags":{},"order":8,"mode":0,"inputs":[{"name":"model","localized_name":"model","type":"PTMODEL","link":47}],"outputs":[{"name":"PTOPTIMIZER","localized_name":"PTOPTIMIZER","type":"PTOPTIMIZER","links":[42],"slot_index":0}],"properties":{"Node name for S&R":"PtoAdamW"},"widgets_values":[0.00005000000000000003,0.9,0.999,0.01,false]},{"id":27,"type":"PtTrainClassificationTransformerModel","pos":[3090.223876953125,1736.2188720703125],"size":[344.3999938964844,292.4398193359375],"flags":{},"order":11,"mode":0,"inputs":[{"name":"model","localized_name":"model","type":"PTMODEL","link":49},{"name":"train_loader","localized_name":"train_loader","type":"PTDATALOADER","link":34},{"name":"optimizer","localized_name":"optimizer","type":"PTOPTIMIZER","link":42},{"name":"loss_function","localized_name":"loss_function","type":"PTLOSS","link":36},{"name":"scheduler","localized_name":"scheduler","type":"PTLRSCHEDULER","shape":7,"link":null},{"name":"val_loader","localized_name":"val_loader","type":"PTDATALOADER","shape":7,"link":35}],"outputs":[{"name":"Model","localized_name":"Model","type":"PTMODEL","links":[38],"slot_index":0},{"name":"train loss","localized_name":"train loss","type":"TENSOR","links":null},{"name":"val loss","localized_name":"val loss","type":"TENSOR","links":null}],"properties":{"Node name for S&R":"PtTrainClassificationTransformerModel"},"widgets_values":[2,true,false,10,false,true]},{"id":34,"type":"MarkdownNote","pos":[1492.9459228515625,2107.386474609375],"size":[286.51580810546875,101.67364501953125],"flags":{},"order":1,"mode":0,"inputs":[],"outputs":[],"properties":{},"widgets_values":["Set max_length to 0 to use model's max_length."],"color":"#432","bgcolor":"#653"},{"id":30,"type":"HfTokenizerEncode","pos":[1482.29150390625,1882.735595703125],"size":[315,154],"flags":{},"order":4,"mode":0,"inputs":[{"name":"model_name","type":"STRING","widget":{"name":"model_name"},"link":44}],"outputs":[{"name":"PTCALLABLE","localized_name":"PTCALLABLE","type":"PTCALLABLE","links":[50,51],"slot_index":0}],"properties":{"Node name for S&R":"HfTokenizerEncode"},"widgets_values":["",true,"max_length",true,0]},{"id":37,"type":"PtvHfLocalDataset","pos":[1910.301025390625,2015.214111328125],"size":[315,178],"flags":{},"order":7,"mode":0,"inputs":[{"name":"encode","localized_name":"encode","type":"PTCALLABLE","link":51}],"outputs":[{"name":"PTVDATASET","localized_name":"PTVDATASET","type":"PTVDATASET","links":[53],"slot_index":0}],"properties":{"Node name for S&R":"PtvHfLocalDataset"},"widgets_values":["imdb/test.jsonl","json","text","label",true,true]},{"id":17,"type":"PtSaveModel","pos":[3495.453125,1735.56494140625],"size":[318.3275146484375,91.27499389648438],"flags":{},"order":12,"mode":0,"inputs":[{"name":"model","localized_name":"model","type":"PTMODEL","link":38}],"outputs":[{"name":"model full path","localized_name":"model full path","type":"STRING","links":[20],"slot_index":0}],"properties":{"Node name for S&R":"PtSaveModel"},"widgets_values":["/container/transfer/models/distilbert_classification_local_dataset.pt"]},{"id":35,"type":"MarkdownNote","pos":[1910.301025390625,2253.2001953125],"size":[304.1402282714844,151.7754669189453],"flags":{},"order":2,"mode":0,"inputs":[],"outputs":[],"properties":{},"widgets_values":["* Make sure to download the IMDB dataset and place them under datasets/imdb using the script save_hf_dataset_locally.py under /dev/scripts before you run this workflow.\n\n* Set encode_return_dict to True for Hugging Face's models."],"color":"#322","bgcolor":"#533"},{"id":36,"type":"PtvHfLocalDataset","pos":[1910.301025390625,1732.4627685546875],"size":[315,178],"flags":{},"order":6,"mode":0,"inputs":[{"name":"encode","localized_name":"encode","type":"PTCALLABLE","link":50}],"outputs":[{"name":"PTVDATASET","localized_name":"PTVDATASET","type":"PTVDATASET","links":[52],"slot_index":0}],"properties":{"Node name for S&R":"PtvHfLocalDataset"},"widgets_values":["imdb/train.jsonl","json","text","label",true,true]},{"id":31,"type":"CDAStringCreate","pos":[1485.653564453125,1684.9422607421875],"size":[315,58],"flags":{},"order":3,"mode":0,"inputs":[],"outputs":[{"name":"STRING","localized_name":"STRING","type":"STRING","links":[44,48],"slot_index":0}],"properties":{"Node name for S&R":"CDAStringCreate"},"widgets_values":["distilbert-base-uncased"]}],"links":[[20,17,0,19,0,"STRING"],[34,7,0,27,1,"PTDATALOADER"],[35,12,0,27,5,"PTDATALOADER"],[36,16,0,27,3,"PTLOSS"],[38,27,0,17,0,"PTMODEL"],[42,29,0,27,2,"PTOPTIMIZER"],[44,31,0,30,0,"STRING"],[47,32,0,29,0,"PTMODEL"],[48,31,0,32,0,"STRING"],[49,32,0,27,0,"PTMODEL"],[50,30,0,36,0,"PTCALLABLE"],[51,30,0,37,0,"PTCALLABLE"],[52,36,0,7,0,"PTVDATASET"],[53,37,0,12,0,"PTVDATASET"]],"groups":[],"config":{},"extra":{"ds":{"scale":1,"offset":[0,0]},"node_versions":{"data-analysis":"68dd090367adffa569c3b3993d56e422f4516d2d","pt-wrapper":"30132509199c6d8ca40c39e39fd37c0e53fd7cdf"},"ue_links":[]},"version":0.4} \ No newline at end of file diff --git a/modules/hugging_face_wrapper/hf_dataset_with_token_encode.py b/modules/hugging_face_wrapper/hf_dataset_with_token_encode.py index c6953d7..4f289ea 100644 --- a/modules/hugging_face_wrapper/hf_dataset_with_token_encode.py +++ b/modules/hugging_face_wrapper/hf_dataset_with_token_encode.py @@ -4,7 +4,7 @@ from html import unescape import torch from torch.utils.data import Dataset from datasets import load_dataset # Hugging Face (not PyTorch) - +from .utils import drop_html_tags class HfDatasetWithTokenEncode(Dataset): """ @@ -54,16 +54,7 @@ class HfDatasetWithTokenEncode(Dataset): # Get a text sample from database and convert to a word embedding sample = self.dataset[idx] if self.remove_html_tags: - # Decode HTML entities like "&", "<", etc. to their actual characters - text = unescape(sample[self.sample_field_name]) - - # Remove HTML tags (like
,

, etc.) - # The regex "<.*?>" uses .*? (non-greedy match) to ensure it matches one tag at a time, - # instead of greedily removing everything between the first '<' and last '>' - text = re.sub(r"<.*?>", " ", text) - - # Replace consecutive white space to a single white space - text = re.sub(r" +", " ", text) + text = drop_html_tags(sample[self.sample_field_name]) else: text = sample[self.sample_field_name] diff --git a/modules/hugging_face_wrapper/hf_local_dataset.py b/modules/hugging_face_wrapper/hf_local_dataset.py new file mode 100644 index 0000000..7eb3d86 --- /dev/null +++ b/modules/hugging_face_wrapper/hf_local_dataset.py @@ -0,0 +1,95 @@ +import os +import sys +from typing import Tuple, Callable, Optional +import re +from html import unescape +import torch +from torch.utils.data import Dataset +from datasets import load_dataset # Hugging Face (not PyTorch) +from .utils import drop_html_tags + +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_dataset_full_path + +class HfLocalDataset(Dataset): + """ + Hugging Face Dataset wrapper to load datasets stored on the local file system. + + pragma: skip_doc + """ + def __init__(self, + file_path: str, + file_format: str, + sample_field_name: str = "text", + label_field_name: str = "label", + encode: Optional[Callable]=None, + remove_html_tags:Optional[bool]=False, + encode_return_dict:Optional[bool]=False) -> None: + """ + Initializes the dataset by loading Hugging Face dataset. + + Args: + file_path (str): Path to the dataset file. Can be an absolute path or a relative path. + If a relative path is provided, it is resolved from the `datasets` directory under this extension. + For example, to load `datasets/foo/bar_train.jsonp`, specify `foo/bar_train.jsonp` or the absolute path. + file_format (str): File format (e.g. json if the file is in jsonl format). + sample_field_name (str): Field name for text samples. + label_field_name (str): Field name for labels. + encode (Optional[Callable]): Token encode function + remove_html_tags (Optional[bool]): Remove html tags in text if True. + encode_return_dict (Optional[bool]): Encode function returns a Dict instead of a Tuple. + """ + dataset_full_path = get_dataset_full_path(file_path) + + # Note split="train" is used here to load the whole dataset + # and not to load just the train dataset. + self.dataset = load_dataset(file_format, + data_files=dataset_full_path, + split="train") + self.sample_field_name = sample_field_name + self.label_field_name = label_field_name + self.encode = encode + self.remove_html_tags = remove_html_tags + self.encode_return_dict = encode_return_dict + + def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Retrieves the (tokens, label) pair at the given index. + + Args: + idx (int): Index of the item. + + Returns: + Tuple[Tuple[torch.Tensor, torch.Tensor], torch.Tensor]: Tuple of the following: + - Tokens tensor and masks tensor. + - Label tensor + """ + # Get a text sample from database and convert to a word embedding + sample = self.dataset[idx] + if self.remove_html_tags: + text = drop_html_tags(sample[self.sample_field_name]) + else: + text = sample[self.sample_field_name] + + if self.encode_return_dict: + output = self.encode(text) + tokens = output["input_ids"] + masks = output["attention_mask"] + else: + tokens, masks = self.encode(text) + + # Get a corresponding label + label = torch.tensor(sample[self.label_field_name], dtype=torch.long) + return (tokens.squeeze(0), masks.squeeze(0)), label.squeeze(0) + + + def __len__(self) -> int: + """ + Returns the number of samples in the dataset. + + Returns: + int: Dataset length. + """ + return len(self.dataset) diff --git a/modules/hugging_face_wrapper/utils.py b/modules/hugging_face_wrapper/utils.py new file mode 100644 index 0000000..7cc1857 --- /dev/null +++ b/modules/hugging_face_wrapper/utils.py @@ -0,0 +1,27 @@ +import re +from html import unescape + + +def drop_html_tags(inputs:str) -> str: + """ + Removes html tags from input. + + Args: + inputs (str): Input text. + + Returns: + str: Text without html tags. + """ + + # Decode HTML entities like "&", "<", etc. to their actual characters + text = unescape(inputs) + + # Remove HTML tags (like
,

, etc.) + # The regex "<.*?>" uses .*? (non-greedy match) to ensure it matches one tag at a time, + # instead of greedily removing everything between the first '<' and last '>' + text = re.sub(r"<.*?>", " ", text) + + # Replace consecutive white space to a single white space + text = re.sub(r" +", " ", text) + + return text \ No newline at end of file diff --git a/modules/pytorch_wrapper/ptv_hf_local_dataset.py b/modules/pytorch_wrapper/ptv_hf_local_dataset.py new file mode 100644 index 0000000..0cabc35 --- /dev/null +++ b/modules/pytorch_wrapper/ptv_hf_local_dataset.py @@ -0,0 +1,90 @@ +import os +import sys +from typing import Any, Dict, Tuple, Callable +import torch + +PROJECT_ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", "..")) +MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules") +sys.path.append(MODULE_ROOT) +from hugging_face_wrapper.hf_local_dataset import HfLocalDataset + +class PtvHfLocalDataset: + """ + A PyTorch Dataset class wrapper for a Hugging Face dataset to load a dataset that is stored on the local file system. + + Args: + file_path (str): Path to the dataset file. Can be an absolute path or a relative path. + If a relative path is provided, it is resolved from the `datasets` directory under this extension. + For example, to load `datasets/foo/bar_train.jsonp`, specify `foo/bar_train.jsonp` or the absolute path. + file_format (str): File format (e.g. json if the file is in jsonl format). + sample_field_name (str): Field name for text samples. + label_field_name (str): Field name for labels. + encode (PTCALLABLE): The reference to a token encoder function. + remove_html_tags (bool): Remove html tags in text if True. + encode_return_dict (bool): Encode function returns a Dict instead of a Tuple. + + 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": { + "file_path": ("STRING", {"default": "", "multiline": False}), + "file_format": ("STRING", {"default": "json","multiline": False}), + "sample_field_name": ("STRING", {"default": "text","multiline": False}), + "label_field_name": ("STRING", {"default": "label","multiline": False}), + "encode": ("PTCALLABLE", {}), + "remove_html_tags": ("BOOLEAN", {"default": False}), + "encode_return_dict": ("BOOLEAN", {"default": False}), + } + } + + RETURN_TYPES: Tuple[str] = ("PTVDATASET",) + FUNCTION: str = "f" + CATEGORY: str = "Training" + + @classmethod + def IS_CHANGED(cls, **kw): + return float("NaN") + + def f(self, + file_path: str, + file_format: str, + sample_field_name: str, + label_field_name: str, + encode: Callable, + remove_html_tags:bool, + encode_return_dict:bool) -> Tuple: + """ + Loads a dataset from Hugging Face with specified parameters. + + Args: + file_path (str): Path to the dataset file. Can be an absolute path or a relative path. + If a relative path is provided, it is resolved from the `datasets` directory under this extension. + For example, to load `datasets/foo/bar_train.jsonp`, specify `foo/bar_train.jsonp` or the absolute path. + file_format (str): File format (e.g. json if the file is in jsonl format). sample_field_name (str): Field name for text samples. + label_field_name (str): Field name for labels. + encode (Callable): The reference to a token encoder function. + remove_html_tags (bool): Remove html tags in text if True. + encode_return_dict (bool): Encode function returns a Dict instead of a Tuple. + + Returns: + Tuple: A tuple containing the dataset instance. + """ + with torch.inference_mode(False): + dc = HfLocalDataset( + file_path=file_path, + file_format=file_format, + sample_field_name=sample_field_name, + label_field_name=label_field_name, + encode=encode, + remove_html_tags=remove_html_tags, + encode_return_dict=encode_return_dict) + return (dc,) diff --git a/pyproject.toml b/pyproject.toml index 0d90451..a1736ff 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.6.0" +version = "3.6.1" license = {file = "LICENSE"} [project.urls] diff --git a/tests/hugging_face_wrapper/hf_local_dataset_with_tokenize_with_hf_tokenize_encode_test.py b/tests/hugging_face_wrapper/hf_local_dataset_with_tokenize_with_hf_tokenize_encode_test.py new file mode 100644 index 0000000..073897e --- /dev/null +++ b/tests/hugging_face_wrapper/hf_local_dataset_with_tokenize_with_hf_tokenize_encode_test.py @@ -0,0 +1,63 @@ +import os +import sys +import unittest +import torch + +PROJECT_ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", "..")) +MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules") +sys.path.append(MODULE_ROOT) + +from hugging_face_wrapper.hf_local_dataset import HfLocalDataset +from hugging_face_wrapper.hf_tokenizer_encode import HfTokenizerEncode + + +class TestHfLocalDataset(unittest.TestCase): + """ + Tests Hugging Face Dataset wrapper with custom token_encode function. + """ + + def setUp(self): + """Set up test instance.""" + self.max_length=15 + encode = HfTokenizerEncode().f( + model_name="distilbert-base-uncased", + padding=True, + padding_method="longest", + truncation=True, + max_length=self.max_length + )[0] + + self.ds = HfLocalDataset( + "imdb/train.jsonl", # dataset file path + "json", + encode=encode, + encode_return_dict=True + ) + + + def test_loading(self): + + # Check number of examples + self.assertEqual(25000, len(self.ds)) + + 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/hugging_face_wrapper/utils_test.py b/tests/hugging_face_wrapper/utils_test.py new file mode 100644 index 0000000..b4f754f --- /dev/null +++ b/tests/hugging_face_wrapper/utils_test.py @@ -0,0 +1,29 @@ +import os +import sys +import unittest +import torch +from itertools import product + +PROJECT_ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", "..")) +MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules") +sys.path.append(MODULE_ROOT) + +from hugging_face_wrapper.utils import drop_html_tags + + +class TestUtils(unittest.TestCase): + def setUp(self): + self.drop_html_tags_data = [ + ("Hello", "Hello"), + ("Hello,
World", "Hello, World"), + ("Hello, test World", "Hello, test World"), + ] + + def test_drop_html_tags(self): + for example, expected in self.drop_html_tags_data: + actual = drop_html_tags(example) + self.assertEqual(expected, actual, f"expected {expected} and actual {actual} do not match.") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/pytorch_wrapper/ptv_hf_local_dataset_test.py b/tests/pytorch_wrapper/ptv_hf_local_dataset_test.py new file mode 100644 index 0000000..d644ee7 --- /dev/null +++ b/tests/pytorch_wrapper/ptv_hf_local_dataset_test.py @@ -0,0 +1,68 @@ +import os +import sys +import unittest +import torch + +PROJECT_ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", "..")) +MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules") +sys.path.append(MODULE_ROOT) + +from hugging_face_wrapper.hf_tokenizer_encode import HfTokenizerEncode +from pytorch_wrapper.ptv_hf_local_dataset import PtvHfLocalDataset + + +class TestPtvHfLocalDataset(unittest.TestCase): + """ + Tests the node for Hugging Face Dataset wrapper with word embedding support. + """ + + def setUp(self): + """Set up test instance.""" + + self.max_length=15 + + encode = HfTokenizerEncode().f( + model_name="distilbert-base-uncased", + padding=True, + padding_method="longest", + truncation=True, + max_length=self.max_length + )[0] + + self.ds = PtvHfLocalDataset().f( + "imdb/train.jsonl", # dataset file path + "json", + sample_field_name = "text", + label_field_name = "label", + encode=encode, + encode_return_dict=True, + remove_html_tags=True + )[0] + + def test_loading(self): + + # Check number of examples + self.assertEquals(25000, len(self.ds)) + + 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()