135 lines
5.0 KiB
Python
135 lines
5.0 KiB
Python
from typing import Any, Dict, Union, List, Tuple, Callable
|
|
import torch
|
|
from transformers import AutoTokenizer
|
|
|
|
|
|
class HfTokenizerEncode:
|
|
"""
|
|
Hugging Face tokenizer wrapper for encoding text into token ID tensors.
|
|
|
|
Args:
|
|
model_name (str): Name of the 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`.
|
|
truncation (bool): Whether to truncate sequences to `max_length`.
|
|
max_length (int): The target length for truncation or padding.
|
|
If you specify 0, it uses model's maximum input length.
|
|
|
|
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": {
|
|
"model_name": ("STRING", {"default":"", "multiline":False}),
|
|
"padding": ("BOOLEAN", {"default": True}),
|
|
"padding_method": (("max_length", "longest"),),
|
|
"truncation": ("BOOLEAN", {"default": True}),
|
|
"max_length": ("INT", {"default": 512, "min": 0, "max": 1_000_000}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES: Tuple[str] = ("PTCALLABLE",)
|
|
FUNCTION: str = "f"
|
|
CATEGORY: str = "Data Analysis"
|
|
|
|
def encode(
|
|
self,
|
|
sentence: Union[str, List[str]],
|
|
) -> Dict[str, torch.Tensor]:
|
|
|
|
"""
|
|
Encodes a string or list of strings into token ID tensors.
|
|
|
|
Args:
|
|
sentence (Union[str, List[str]]): Input sentence(s).
|
|
|
|
Returns:
|
|
Dict:
|
|
input_ids: token IDs
|
|
attention_mask: Masks. 1 is non-padding, 0 is padding.
|
|
"""
|
|
if self.padding:
|
|
if self.max_length == 0:
|
|
retval = self.tokenizer(
|
|
sentence,
|
|
padding=self.padding_method,
|
|
truncation=self.truncation,
|
|
return_tensors="pt"
|
|
)
|
|
else:
|
|
retval = self.tokenizer(
|
|
sentence,
|
|
padding=self.padding_method,
|
|
truncation=self.truncation,
|
|
max_length=self.max_length,
|
|
return_tensors="pt"
|
|
)
|
|
else: # no padding
|
|
if self.max_length == 0:
|
|
retval = self.tokenizer(
|
|
sentence,
|
|
padding=False,
|
|
truncation=self.truncation,
|
|
return_tensors="pt"
|
|
)
|
|
else:
|
|
retval = self.tokenizer(
|
|
sentence,
|
|
padding=False,
|
|
truncation=self.truncation,
|
|
return_tensors="pt"
|
|
)
|
|
return retval
|
|
|
|
def f(
|
|
self,
|
|
model_name: str,
|
|
padding: bool,
|
|
padding_method: str,
|
|
truncation: bool,
|
|
max_length: int
|
|
) -> Tuple[Callable[[Union[str, List[str]]], Tuple[Union[torch.Tensor, List[torch.Tensor]], Union[torch.Tensor, List[torch.Tensor]]]]]:
|
|
"""
|
|
Instantiates a Hugging Face tokenizer, binds configuration and returns a callable encoder function.
|
|
- Supports includes:
|
|
- BertTokenizer
|
|
- RobertaTokenizer
|
|
- DistilBertTokenizer
|
|
- AlbertTokenizer
|
|
- ElectraTokenizer
|
|
|
|
Args:
|
|
model_name (str): Name of the 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`.
|
|
truncation (bool): Whether to truncate sequences to `max_length`.
|
|
max_length (int): The target length for truncation or padding.
|
|
If you specify 0, it uses model's maximum input length.
|
|
|
|
Returns:
|
|
Tuple[Callable]: A tuple containing a callable for the tokenizer function.
|
|
"""
|
|
|
|
self.tokenizer = AutoTokenizer.from_pretrained(model_name)
|
|
self.padding = padding
|
|
self.padding_method = padding_method
|
|
self.truncation = truncation
|
|
self.max_length = max_length
|
|
return (self.encode,)
|