115 lines
4.5 KiB
Python
115 lines
4.5 KiB
Python
# These codes are copied from modelscope revision c58451baead80d83281f063d12fb377fad415257
|
|
# Copyright (c) Alibaba, Inc. and its affiliates.
|
|
|
|
import os
|
|
from collections.abc import Mapping
|
|
|
|
import json
|
|
from transformers import AutoTokenizer
|
|
|
|
from modelscope.metainfo import Models
|
|
from modelscope.outputs import OutputKeys
|
|
from modelscope.utils.constant import ModeKeys
|
|
from modelscope.utils.logger import get_logger
|
|
|
|
logger = get_logger()
|
|
|
|
__all__ = [
|
|
'NLPTokenizer',
|
|
]
|
|
|
|
|
|
class NLPTokenizer:
|
|
|
|
def __init__(self,
|
|
model_dir: str = None,
|
|
model_type=None,
|
|
use_fast: bool = None,
|
|
tokenize_kwargs=None):
|
|
"""The transformers tokenizer preprocessor base class.
|
|
|
|
Any nlp preprocessor which uses the huggingface tokenizer can inherit from this class.
|
|
|
|
Args:
|
|
model_dir (str, `optional`): The local path containing the files used to create a preprocessor.
|
|
use_fast (str, `optional`): Use the fast version of tokenizer
|
|
tokenize_kwargs (dict, `optional`): These args will be directly fed into the tokenizer.
|
|
"""
|
|
self.model_dir = model_dir
|
|
self.model_type = model_type
|
|
self.tokenize_kwargs = tokenize_kwargs
|
|
if self.tokenize_kwargs is None:
|
|
self.tokenize_kwargs = {}
|
|
self._use_fast = use_fast
|
|
self._tokenizer = None
|
|
|
|
@property
|
|
def tokenizer(self):
|
|
if self._tokenizer is None:
|
|
self._tokenizer = self.build_tokenizer()
|
|
return self._tokenizer
|
|
|
|
@property
|
|
def use_fast(self):
|
|
if self._use_fast is None:
|
|
if self._use_fast is None and self.model_dir is None:
|
|
self._use_fast = False
|
|
elif self._use_fast is None and os.path.isfile(
|
|
os.path.join(self.model_dir, 'tokenizer_config.json')):
|
|
with open(
|
|
os.path.join(self.model_dir, 'tokenizer_config.json'),
|
|
'r',
|
|
encoding='utf-8') as f:
|
|
json_config = json.load(f)
|
|
self._use_fast = json_config.get('use_fast')
|
|
self._use_fast = False if self._use_fast is None else self._use_fast
|
|
return self._use_fast
|
|
|
|
def build_tokenizer(self):
|
|
"""Build a tokenizer by the model type.
|
|
|
|
NOTE: The fast tokenizers have a multi-thread problem, use it carefully.
|
|
|
|
Returns:
|
|
The initialized tokenizer.
|
|
"""
|
|
# fast version lead to parallel inference failed
|
|
model_type = self.model_type
|
|
model_dir = self.model_dir
|
|
|
|
if model_type in (Models.structbert, Models.gpt3, Models.palm,
|
|
Models.plug, Models.megatron_bert,
|
|
Models.plug_mental, Models.fid_plug):
|
|
from transformers import BertTokenizer, BertTokenizerFast
|
|
tokenizer = BertTokenizerFast if self.use_fast else BertTokenizer
|
|
return tokenizer.from_pretrained(
|
|
model_dir) if model_dir is not None else tokenizer()
|
|
elif model_type == Models.veco:
|
|
from transformers import XLMRobertaTokenizer, XLMRobertaTokenizerFast
|
|
tokenizer = XLMRobertaTokenizerFast if self.use_fast else XLMRobertaTokenizer
|
|
return tokenizer.from_pretrained(
|
|
model_dir) if model_dir is not None else tokenizer()
|
|
elif model_type == Models.llama:
|
|
from modelscope.models.nlp import LlamaTokenizer, LlamaTokenizerFast
|
|
tokenizer = LlamaTokenizerFast if self.use_fast else LlamaTokenizer
|
|
return tokenizer.from_pretrained(
|
|
model_dir) if model_dir is not None else tokenizer()
|
|
|
|
assert model_dir is not None
|
|
return AutoTokenizer.from_pretrained(model_dir, use_fast=self.use_fast)
|
|
|
|
def __call__(self, text, text_pair=None, **kwargs):
|
|
kwargs['max_length'] = kwargs.get('max_length',
|
|
kwargs.pop('sequence_length', None))
|
|
if kwargs['max_length'] is None:
|
|
kwargs.pop('max_length')
|
|
tokenize_kwargs = {k: v for k, v in self.tokenize_kwargs.items()}
|
|
tokenize_kwargs.update(kwargs)
|
|
kwargs.update(self.tokenize_kwargs)
|
|
return self.tokenizer(text, text_pair, **tokenize_kwargs)
|
|
|
|
def get_tokenizer_kwarg(self, key, default_value=None):
|
|
if key in self.tokenize_kwargs:
|
|
return self.tokenize_kwargs[key]
|
|
return self.tokenizer.init_kwargs.get(key, default_value)
|