hankcs--hanlp
wehub-resource-sync
82 行
4.2 KiB
Python
82 行
4.2 KiB
Python
# -*- coding:utf-8 -*-
|
|
# Author: hankcs
|
|
# Date: 2021-05-13 13:24
|
|
from typing import Union, Optional
|
|
|
|
from transformers import BertTokenizerFast, TensorType, BatchEncoding, BertJapaneseTokenizer as _BertJapaneseTokenizer
|
|
from transformers.file_utils import PaddingStrategy
|
|
from transformers.tokenization_utils_base import TextInput, PreTokenizedInput, EncodedInput, TruncationStrategy
|
|
|
|
|
|
class BertJapaneseTokenizer(_BertJapaneseTokenizer):
|
|
# We may need to customize character level tokenization to handle English words and URLs
|
|
pass
|
|
|
|
|
|
class BertJapaneseTokenizerFast(BertTokenizerFast):
|
|
def encode_plus(
|
|
self,
|
|
text: Union[TextInput, PreTokenizedInput, EncodedInput],
|
|
text_pair: Optional[Union[TextInput, PreTokenizedInput, EncodedInput]] = None,
|
|
add_special_tokens: bool = True,
|
|
padding: Union[bool, str, PaddingStrategy] = False,
|
|
truncation: Union[bool, str, TruncationStrategy] = False,
|
|
max_length: Optional[int] = None,
|
|
stride: int = 0,
|
|
is_split_into_words: bool = False,
|
|
pad_to_multiple_of: Optional[int] = None,
|
|
return_tensors: Optional[Union[str, TensorType]] = None,
|
|
return_token_type_ids: Optional[bool] = None,
|
|
return_attention_mask: Optional[bool] = None,
|
|
return_overflowing_tokens: bool = False,
|
|
return_special_tokens_mask: bool = False,
|
|
return_offsets_mapping: bool = False,
|
|
return_length: bool = False,
|
|
verbose: bool = True,
|
|
**kwargs
|
|
) -> BatchEncoding:
|
|
"""
|
|
Tokenize and prepare for the model a sequence or a pair of sequences.
|
|
|
|
.. warning::
|
|
This method is deprecated, ``__call__`` should be used instead.
|
|
|
|
Args:
|
|
text (:obj:`str`, :obj:`List[str]` or :obj:`List[int]` (the latter only for not-fast tokenizers)):
|
|
The first sequence to be encoded. This can be a string, a list of strings (tokenized string using the
|
|
``tokenize`` method) or a list of integers (tokenized string ids using the ``convert_tokens_to_ids``
|
|
method).
|
|
text_pair (:obj:`str`, :obj:`List[str]` or :obj:`List[int]`, `optional`):
|
|
Optional second sequence to be encoded. This can be a string, a list of strings (tokenized string using
|
|
the ``tokenize`` method) or a list of integers (tokenized string ids using the
|
|
``convert_tokens_to_ids`` method).
|
|
"""
|
|
text = list(text)
|
|
is_split_into_words = True
|
|
encoding = BertJapaneseTokenizer.encode_plus(self,
|
|
text,
|
|
text_pair,
|
|
add_special_tokens,
|
|
padding,
|
|
truncation,
|
|
max_length,
|
|
stride,
|
|
is_split_into_words,
|
|
pad_to_multiple_of,
|
|
return_tensors,
|
|
return_token_type_ids,
|
|
return_attention_mask,
|
|
return_overflowing_tokens,
|
|
return_special_tokens_mask,
|
|
return_offsets_mapping,
|
|
return_length,
|
|
verbose,
|
|
**kwargs
|
|
)
|
|
offsets = encoding.encodings[0].offsets
|
|
fixed_offsets = [(b + i, e + i) for i, (b, e) in enumerate(offsets)]
|
|
# TODO: This doesn't work with rust tokenizers
|
|
encoding.encodings[0].offsets.clear()
|
|
encoding.encodings[0].offsets.extend(fixed_offsets)
|
|
return encoding
|