135 lines
5.0 KiB
Python
135 lines
5.0 KiB
Python
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,)
|