, 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()