81 lines
2.9 KiB
Python
81 lines
2.9 KiB
Python
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
|
|
|
|
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: 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:
|
|
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 (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.
|
|
"""
|
|
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
|
|
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)
|