diff --git a/dooly/build_utils.py b/dooly/build_utils.py deleted file mode 100644 index a2424dc..0000000 --- a/dooly/build_utils.py +++ /dev/null @@ -1,66 +0,0 @@ -from typing import Dict, Union, Optional - -import torch - -from transformers.file_utils import hf_bucket_url, cached_path - - -HUB_NAME = "jinmang2/dooly-hub" -VOCAB_NAME = "vocab.json" -CONFIG_NAME = "config.json" -WEIGHTS_NAME = "pytorch_model.bin" - -TOKENIZER_USER_AGENT = { - "file_type": "tokenizer", - "from_auto_class": False, - "is_fast": False, -} - -CONFIG_USER_AGENT = { - "file_type": "config", - "from_auto_class": False, -} - -MODEL_USER_AGENT = { - "file_type": "model", - "framework": "pytorch", - "from_auto_class": False, -} - -# @TODO: cuda:0, 1, 2 -DEFAULT_DEVICE = "cuda" if torch.cuda.is_available() else "cpu" - - -def download_from_hf_hub( - model_id: str, - filename: str, - subfolder: Optional[str] = None, - revision: Optional[str] = None, - cache_dir: Optional[str] = None, - force_download: bool = False, - resume_download: bool = False, - user_agent: Optional[Union[Dict, str]] = None, -) -> Optional[str]: - # Resolve a model identifier, a file name, and an optional revision id, - # to a huggingface.co-hosted url, redirecting to Cloudfront - # (a Content Delivery Network, or CDN) for large files. - huggingface_co_resolved_file = hf_bucket_url( - model_id=model_id, - filename=filename, - subfolder=subfolder, - revision=revision, - ) - # Given something that might be a URL (or might be a local path), - # determine which. If it's a URL, download the file and cache it, - # and return the path to the cached file. If it's already a local path, - # make sure the file exists and then return the path - # Do not extract files (extract_compressed_file and force_extract is False) - resolved_file_path = cached_path( - huggingface_co_resolved_file, - cache_dir=cache_dir, - force_download=force_download, - resume_download=resume_download, - user_agent=user_agent, - ) - - return resolved_file_path diff --git a/dooly/converters/base.py b/dooly/converters/base.py index b374bdf..f206671 100644 --- a/dooly/converters/base.py +++ b/dooly/converters/base.py @@ -70,10 +70,7 @@ def convert(self): config = self.get_model_config(self._pororo_model) hf_model = self.intialize_hf_model(config) - self._hf_model = self.porting_pororo_to_hf( - self._pororo_model, - hf_model, - ) + self._hf_model = self.porting_pororo_to_hf(self._pororo_model, hf_model) self.save_hf_model(self._hf_model) # convert misc files @@ -116,7 +113,8 @@ def load_vocab(self): raise NotImplementedError def save_vocab(self, vocab, filename: str = "vocab.json"): - with open(os.path.join(self.save_path, filename), "w", encoding="utf-8") as f: + save_file_path = os.path.join(self.save_path, filename) + with open(save_file_path, "w", encoding="utf-8") as f: json.dump(vocab, f, ensure_ascii=False) def load_and_save_vocab(self): @@ -141,9 +139,9 @@ def get_misc_filenames(self): def load_and_save_misc(self): misc_files = self.get_misc_filenames() for misc_file in misc_files: - shutil.move( - src=os.path.join(self._pororo_save_path, misc_file), dst=self.save_path - ) + src_path = os.path.join(self._pororo_save_path, misc_file) + dst_path = self.save_path + shutil.move(src=src_path, dst=dst_path) class FsmtConverter(DoolyConverter): @@ -260,19 +258,14 @@ def porting_pororo_to_hf(self, pororo_model, hf_model): sent_encoder = pororo_model.model.encoder.sentence_encoder # Now let's copy all the weights. # Embeddings - hf_model.roberta.embeddings.word_embeddings.weight = ( - sent_encoder.embed_tokens.weight - ) - hf_model.roberta.embeddings.position_embeddings.weight = ( - sent_encoder.embed_positions.weight - ) - hf_model.roberta.embeddings.token_type_embeddings.weight.data = ( - torch.zeros_like(hf_model.roberta.embeddings.token_type_embeddings.weight) + embeddings = hf_model.roberta.embeddings + embeddings.word_embeddings.weight = sent_encoder.embed_tokens.weight + embeddings.position_embeddings.weight = sent_encoder.embed_positions.weight + embeddings.token_type_embeddings.weight.data = torch.zeros_like( + hf_model.roberta.embeddings.token_type_embeddings.weight ) # just zero them out b/c RoBERTa doesn't use them. - hf_model.roberta.embeddings.LayerNorm.weight = ( - sent_encoder.emb_layer_norm.weight - ) - hf_model.roberta.embeddings.LayerNorm.bias = sent_encoder.emb_layer_norm.bias + embeddings.LayerNorm.weight = sent_encoder.emb_layer_norm.weight + embeddings.LayerNorm.bias = sent_encoder.emb_layer_norm.bias for i in range(hf_model.config.num_hidden_layers): # Encoder: start of layer diff --git a/dooly/converters/kobart_utils.py b/dooly/converters/kobart_utils.py index 6e14c71..073d7a5 100644 --- a/dooly/converters/kobart_utils.py +++ b/dooly/converters/kobart_utils.py @@ -17,11 +17,7 @@ def is_available_boto3(): class AwsS3Downloader(object): - def __init__( - self, - aws_access_key_id=None, - aws_secret_access_key=None, - ): + def __init__(self, aws_access_key_id=None, aws_secret_access_key=None): self.resource = boto3.Session( aws_access_key_id=aws_access_key_id, aws_secret_access_key=aws_secret_access_key, diff --git a/dooly/converters/task_specific.py b/dooly/converters/task_specific.py index 4ef7f08..6457458 100644 --- a/dooly/converters/task_specific.py +++ b/dooly/converters/task_specific.py @@ -6,7 +6,7 @@ from .base import DoolyConverter, FsmtConverter, RobertaConverter from .kobart_utils import download -from ..tokenizers.hf_tokenizer import ( +from ..tokenizers.fast import ( convert_vocab_from_fairseq_to_hf, build_custom_roberta_tokenizer, PreTrainedTokenizerFast, diff --git a/dooly/dooly.py b/dooly/dooly.py index 6643502..375222b 100644 --- a/dooly/dooly.py +++ b/dooly/dooly.py @@ -2,6 +2,7 @@ from typing import Optional from .tasks import DoolyTaskHub +from .utils import _locate def normalize_task(task: str): @@ -55,7 +56,7 @@ def __new__( raise KeyError( f"Unavailable task name '{task}'. See here {TASK_ALIASES.keys()}" ) - task_cls = DoolyTaskHub[task] + task_cls = _locate(DoolyTaskHub[task]) if lang is not None: lang = LANG_ALIASES.get(lang.lower(), None) return task_cls.build(lang, n_model, **kwargs) @@ -76,9 +77,11 @@ def available_tasks() -> str: def available_models(task: str) -> str: if task not in TASK_ALIASES: raise KeyError( - f"Unknown task {task}. Please check available models via `available_tasks()`." + f"Unknown task {task}. " + "Please check available models via `available_tasks()`." ) - task_cls = DoolyTaskHub[TASK_ALIASES[normalize_task(task)]] + task_cls_path = DoolyTaskHub[TASK_ALIASES[normalize_task(task)]] + task_cls = _locate(task_cls_path) output = f"Available models for `{task}` are " for lang, models in task_cls.available_models.items(): diff --git a/dooly/models/__init__.py b/dooly/models/__init__.py index 974a8cd..940f166 100644 --- a/dooly/models/__init__.py +++ b/dooly/models/__init__.py @@ -1,227 +1,2 @@ -import json -from packaging import version -from contextlib import contextmanager -from typing import Optional - -import torch - -import transformers -from transformers.configuration_utils import PretrainedConfig -from transformers.modeling_utils import PreTrainedModel - -from .modeling_bart import BartForConditionalGeneration -from .modeling_fsmt import FSMTForConditionalGeneration -from .modeling_roberta import ( - RobertaForDependencyParsing, - RobertaForSpanPrediction, - RobertaForSequenceTagging, - RobertaForSequenceClassification, -) -from ..build_utils import ( - download_from_hf_hub, - CONFIG_USER_AGENT, - HUB_NAME, - MODEL_USER_AGENT, - CONFIG_NAME, - WEIGHTS_NAME, -) - - -DoolyModelHub = { - "dp": { - "ko": {"posbert.base": RobertaForDependencyParsing}, - }, - "mrc": { - "ko": {"brainbert.base": RobertaForSpanPrediction}, - }, - "mt": { - "multi": { - "transformer.large.mtpg": FSMTForConditionalGeneration, - "transformer.large.fast.mtpg": FSMTForConditionalGeneration, - }, - }, - "ner": { - "ko": {"charbert.base": RobertaForSequenceTagging}, - "en": {"roberta.base": RobertaForSequenceTagging}, - "ja": {"jaberta.base": RobertaForSequenceTagging}, - "zh": {"zhberta.base": RobertaForSequenceTagging}, - }, - "nli": { - "ko": {"brainbert.base": RobertaForSequenceClassification}, - "en": {"roberta.base": RobertaForSequenceClassification}, - "ja": {"jaberta.base": RobertaForSequenceClassification}, - "zh": {"zhberta.base": RobertaForSequenceClassification}, - }, - "qg": { - "ko": {"kobart.base": BartForConditionalGeneration}, - }, - "wsd": { - "ko": {"transformer.large": FSMTForConditionalGeneration}, - }, -} -DoolyModelHub["bt"] = DoolyModelHub["mt"] -DoolyModelHub["zero_topic"] = DoolyModelHub["nli"] -available_tasks = list(DoolyModelHub.keys()) - -_init_weights = True - - -@contextmanager -def no_init_weights(_enable=True): - global _init_weights - if _enable: - _init_weights = False - try: - yield - finally: - _init_weights = True - - -class DoolyModel: - """ Dooly Model """ - - @classmethod - def build_model(cls, task: str, lang: str, n_model: str, **kwargs): - assert ( - task in available_tasks - ), f"Task `{task}` is not available. See here {available_tasks}." - available_langs = DoolyModelHub[task] - assert lang in available_langs, ( - f"Language `{lang}` is not available in this task {task}. " - f"See here {available_langs}." - ) - available_models = available_langs[lang] - assert n_model in available_models, ( - f"Model `{n_model}` is not available in this task-lang pair. " - f"See here {available_models}." - ) - - model_class = available_models[n_model] - - return cls._build_model(task, lang, n_model, model_class, **kwargs) - - @classmethod - def _build_model_config( - cls, - task: str, - lang: str, - n_model: str, - config_class: PretrainedConfig, - revision: Optional[str] = None, - cache_dir: Optional[str] = None, - force_download: bool = False, - resume_download: bool = False, - **kwargs, - ) -> PretrainedConfig: - # Load from URL or cache if already cached - resolved_config_file = download_from_hf_hub( - model_id=HUB_NAME, - filename=CONFIG_NAME, - subfolder=f"{task}/{lang}/{n_model}", - revision=revision, - cache_dir=cache_dir, - force_download=force_download, - resume_download=resume_download, - user_agent=CONFIG_USER_AGENT, - ) - - # _dict_from_json_file - with open(resolved_config_file, "r", encoding="utf-8") as reader: - text = reader.read() - config_dict = json.loads(text) - - return config_class.from_dict(config_dict, **kwargs) - - @classmethod - def _build_model( - cls, - task: str, - lang: str, - n_model: str, - model_class: PreTrainedModel, - revision: Optional[str] = None, - cache_dir: Optional[str] = None, - force_download: bool = False, - resume_download: bool = False, - low_cpu_mem_usage: bool = False, - _fast_init: bool = True, - **kwargs, - ) -> PreTrainedModel: - if low_cpu_mem_usage: - assert version.parse(torch.__version__) > version.parse("1.9"), ( - "torch>=1.9 is required for a normal functioning of this module" - f"using the low_cpu_mem_usage=={low_cpu_mem_usage}, " - f"but found torch=={torch.__version__}" - ) - - config_class: PretrainedConfig = model_class.config_class - - config = cls._build_model_config( - task=task, - lang=lang, - n_model=n_model, - config_class=config_class, - revision=revision, - cache_dir=cache_dir, - force_download=force_download, - resume_download=resume_download, - **kwargs, - ) - - # Load from URL or cache if already cached - resolved_archive_file = download_from_hf_hub( - model_id=HUB_NAME, - filename=WEIGHTS_NAME, - subfolder=f"{task}/{lang}/{n_model}", - revision=revision, - cache_dir=cache_dir, - force_download=force_download, - resume_download=resume_download, - user_agent=MODEL_USER_AGENT, - ) - state_dict = torch.load(resolved_archive_file, map_location="cpu") - - if low_cpu_mem_usage: - loaded_state_dict_keys = [k for k in state_dict.keys()] - del state_dict # free CPU memory - will reload again later - - with no_init_weights(_enable=_fast_init): - model = model_class(config, **kwargs) - - # There was an update to the from_pretrained method of models in v4.18.0. - # See fetch below. - # ref. https://github.com/huggingface/transformers/releases/tag/v4.18.0 - # ref. https://github.com/huggingface/transformers/pull/16343 - if low_cpu_mem_usage: - kwargs = dict( - model=model, - loaded_state_dict_keys=loaded_state_dict_keys, - resolved_archive_file=resolved_archive_file, - ) - if version.parse(transformers.__version__) >= version.parse("4.18.0"): - load_pretrained_model = model_class._load_pretrained_model_low_mem - else: - load_pretrained_model = model_class._load_state_dict_into_model_low_mem - load_pretrained_model(**kwargs) - else: - kwargs = dict( - model=model, - state_dict=state_dict, - pretrained_model_name_or_path=HUB_NAME, - ignore_mismatched_sizes=False, - _fast_init=_fast_init, - ) - if version.parse(transformers.__version__) >= version.parse("4.18.0"): - kwargs.update(dict(resolved_archive_file=resolved_archive_file)) - load_pretrained_model = model_class._load_pretrained_model - else: - load_pretrained_model = model_class._load_state_dict_into_model - model, _, _, _, _ = load_pretrained_model(**kwargs) - - # make sure token embedding weights are still tied if needed - model.tie_weights() - - # Set model in evaluation mode to deactivate DropOut modules by default - model.eval() - - return model +from .load import DoolyModelHub # noqa +from .load import load_dooly_model # noqa diff --git a/dooly/models/load.py b/dooly/models/load.py new file mode 100644 index 0000000..cbf17a3 --- /dev/null +++ b/dooly/models/load.py @@ -0,0 +1,112 @@ +from typing import Type, Union +import transformers + +from ..utils import _locate, register_subfolder, DOOLY_HUB_NAME + + +DoolyModelHub = { + "dp": {"ko": {"posbert.base": "modeling_roberta.RobertaForDependencyParsing"}}, + "mrc": {"ko": {"brainbert.base": "modeling_roberta.RobertaForSpanPrediction"}}, + "mt": { + "multi": { + "transformer.large.mtpg": "modeling_fsmt.FSMTForConditionalGeneration", + "transformer.large.fast.mtpg": "modeling_fsmt.FSMTForConditionalGeneration", + }, + }, + "ner": { + "ko": {"charbert.base": "modeling_roberta.RobertaForSequenceTagging"}, + "en": {"roberta.base": "modeling_roberta.RobertaForSequenceTagging"}, + "ja": {"jaberta.base": "modeling_roberta.RobertaForSequenceTagging"}, + "zh": {"zhberta.base": "modeling_roberta.RobertaForSequenceTagging"}, + }, + "nli": { + "ko": {"brainbert.base": "modeling_roberta.RobertaForSequenceClassification"}, + "en": {"roberta.base": "modeling_roberta.RobertaForSequenceClassification"}, + "ja": {"jaberta.base": "modeling_roberta.RobertaForSequenceClassification"}, + "zh": {"zhberta.base": "modeling_roberta.RobertaForSequenceClassification"}, + }, + "qg": {"ko": {"kobart.base": "modeling_bart.BartForConditionalGeneration"}}, + "wsd": {"ko": {"transformer.large": "modeling_fsmt.FSMTForConditionalGeneration"}}, +} +DoolyModelHub["bt"] = DoolyModelHub["mt"] +DoolyModelHub["zero_topic"] = DoolyModelHub["nli"] + +available_tasks = list(DoolyModelHub.keys()) + + +def load_pretrained_model( + pretrained_model_name_or_path: str, + model_class: Type[transformers.PreTrainedModel], + **kwargs, +) -> transformers.PreTrainedModel: + return model_class.from_pretrained(pretrained_model_name_or_path, **kwargs) + + +def load_model_from_dooly_hub( + subfolder: str, model_class: Type[transformers.PreTrainedModel], **kwargs +) -> transformers.PreTrainedModel: + @register_subfolder + def _load_pretrained( + pretrained_model_name_or_path: str, subfolder: str, **kwargs + ) -> transformers.PreTrainedModel: + return model_class.from_pretrained(pretrained_model_name_or_path, **kwargs) + + return _load_pretrained( + pretrained_model_name_or_path=DOOLY_HUB_NAME, subfolder=subfolder, **kwargs + ) + + +def load_dooly_model( + pretrained_model_name_or_path: str = None, + model_class: Union[str, Type[transformers.PreTrainedModel]] = None, + task: str = None, + lang: str = None, + n_model: str = None, + **kwargs, +) -> transformers.PreTrainedModel: + if pretrained_model_name_or_path is not None: + if model_class is None: + raise ValueError( + "If you are using the personal huggingface.co model, " + "`model_class` parameter is required." + ) + return load_pretrained_model( + pretrained_model_name_or_path, model_class, **kwargs + ) + + if all([task is None and lang is None and n_model is None]): + raise ValueError( + "`task`, `lang`, and `n_model` parameters are required to " + "access the subfolder of dooly-hub.\nCheck your parameters! " + f"`task`: {task} `lang`: {lang} `n_model`: {n_model}." + ) + + assert ( + task in available_tasks + ), f"Task `{task}` is not available. See here {available_tasks}." + + available_langs = DoolyModelHub[task] + assert lang in available_langs, ( + f"Language `{lang}` is not available in this task {task}. " + f"See here {available_langs}." + ) + + available_models = available_langs[lang] + assert n_model in available_models, ( + f"Model `{n_model}` is not available in this task-lang pair. " + f"See here {available_models}." + ) + + subfolder = f"{task}/{lang}/{n_model}" + + if model_class is None: + module_path = "dooly.models." + available_models[n_model] + elif isinstance(model_class, str): + module_path = model_class + + if not issubclass(model_class, transformers.PreTrainedModel): + model_class = _locate(module_path) + + return load_model_from_dooly_hub( + subfolder=subfolder, model_class=model_class, **kwargs + ) diff --git a/dooly/models/modeling_fsmt.py b/dooly/models/modeling_fsmt.py index bfc01fa..6e495c6 100644 --- a/dooly/models/modeling_fsmt.py +++ b/dooly/models/modeling_fsmt.py @@ -12,8 +12,6 @@ from transformers.modeling_outputs import ( BaseModelOutput, BaseModelOutputWithPastAndCrossAttentions, - # Seq2SeqLMOutput, - # Seq2SeqModelOutput, ) from transformers.deepspeed import is_deepspeed_zero3_enabled @@ -26,10 +24,7 @@ def invert_mask(attention_mask): class FSMTConfig(FSMTConfig): def __init__( - self, - encoder_pre_layernorm=False, - decoder_pre_layernorm=False, - **kwargs, + self, encoder_pre_layernorm=False, decoder_pre_layernorm=False, **kwargs ): self.encoder_pre_layernorm = encoder_pre_layernorm self.decoder_pre_layernorm = decoder_pre_layernorm diff --git a/dooly/models/modeling_roberta.py b/dooly/models/modeling_roberta.py index 531652b..c24a695 100644 --- a/dooly/models/modeling_roberta.py +++ b/dooly/models/modeling_roberta.py @@ -165,7 +165,8 @@ def __init__(self, config): config, "position_embedding_type", "absolute" ) self.register_buffer( - "position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)) + "position_ids", + torch.arange(config.max_position_embeddings).expand((1, -1)), ) if version.parse(torch.__version__) > version.parse("1.6.0"): self.register_buffer( diff --git a/dooly/models/utils/modeling_heads.py b/dooly/models/utils/modeling_heads.py index 11a99ae..781b6df 100644 --- a/dooly/models/utils/modeling_heads.py +++ b/dooly/models/utils/modeling_heads.py @@ -68,9 +68,7 @@ def __init__(self, config): dropout=classifier_dropout, ) self.head_attn_post = nn.MultiheadAttention( - embed_dim=config.hidden_size, - num_heads=1, - dropout=classifier_dropout, + embed_dim=config.hidden_size, num_heads=1, dropout=classifier_dropout ) def forward(self, features, masks=None, **kwargs): @@ -111,9 +109,7 @@ def __init__(self, config): ) # shared with encoder self.gru = nn.GRU( - input_size=self.hidden_size, - hidden_size=self.hidden_size, - batch_first=True, + input_size=self.hidden_size, hidden_size=self.hidden_size, batch_first=True ) self.gating2id = {"none": 0, "dontcare": 1, "ptr": 2, "yes": 3, "no": 4} diff --git a/dooly/models/utils/modeling_utils.py b/dooly/models/utils/modeling_utils.py index a5de0bb..e2e3cc9 100644 --- a/dooly/models/utils/modeling_utils.py +++ b/dooly/models/utils/modeling_utils.py @@ -2,9 +2,7 @@ def masked_cross_entropy_for_value( - logits: torch.Tensor, - target: torch.Tensor, - pad_idx: int = 0, + logits: torch.Tensor, target: torch.Tensor, pad_idx: int = 0 ) -> torch.Tensor: # loss_gen mask = target.ne(pad_idx) logits_flat = logits.view(-1, logits.size(-1)) diff --git a/dooly/tokenizers/__init__.py b/dooly/tokenizers/__init__.py index d75450d..cad3024 100644 --- a/dooly/tokenizers/__init__.py +++ b/dooly/tokenizers/__init__.py @@ -1,161 +1,2 @@ -import json -from functools import partial -from typing import Union, Optional - -from .base import Tokenizer -from .bpe import Gpt2BpeTokenizer, BpeJaZhTokenizer -from .char import CharS1Tokenizer, CharS2Tokenizer -from .hf_tokenizer import ( - PreTrainedTokenizerBase, - PreTrainedTokenizerFast, - RobertaTokenizerFast, -) -from .pos_tagger import ( - PosDpTokenizer, -) -from ..build_utils import ( - download_from_hf_hub, - HUB_NAME, - VOCAB_NAME, - TOKENIZER_USER_AGENT, -) - - -DoolyTokenizerHub = { - "dp": { - "ko": {"posbert.base": PosDpTokenizer}, - }, - "mrc": { - "ko": {"brainbert.base": RobertaTokenizerFast}, - }, - "mt": { - "multi": { - "transformer.large.mtpg": CharS2Tokenizer, - "transformer.large.fast.mtpg": CharS2Tokenizer, - }, - }, - "ner": { - "ko": {"charbert.base": CharS1Tokenizer}, - "en": {"roberta.base": Gpt2BpeTokenizer}, - "ja": {"jaberta.base": BpeJaZhTokenizer}, - "zh": {"zhberta.base": BpeJaZhTokenizer}, - }, - "nli": { - "ko": {"brainbert.base": RobertaTokenizerFast}, - "en": {"roberta.base": Gpt2BpeTokenizer}, - "ja": {"jaberta.base": BpeJaZhTokenizer}, - "zh": {"zhberta.base": BpeJaZhTokenizer}, - }, - "qg": { - "ko": {"kobart.base": PreTrainedTokenizerFast}, - }, - "wsd": { - "ko": {"transformer.large": CharS2Tokenizer}, - }, -} -DoolyTokenizerHub["bt"] = DoolyTokenizerHub["mt"] -DoolyTokenizerHub["zero_topic"] = DoolyTokenizerHub["nli"] -available_tasks = list(DoolyTokenizerHub.keys()) - - -class DoolyTokenizer: - """ Dooly Tokenizer """ - - @classmethod - def build_tokenizer( - cls, task: str, lang: str, n_model: Optional[str] = None, **kwargs - ): - assert ( - task in available_tasks - ), f"Task `{task}` is not available. See here {available_tasks}." - available_langs = DoolyTokenizerHub[task] - assert lang in available_langs, ( - f"Language `{lang}` is not available in this task {task}. " - f"See here {available_langs}." - ) - available_models = available_langs[lang] - if n_model is None: - n_model = list(available_models.keys())[0] - assert n_model in available_models, ( - f"Model `{n_model}` is not available in this task-lang pair. " - f"See here {available_models}." - ) - - tokenizer_class = available_models[n_model] - - return cls._build_tokenizer(task, lang, n_model, tokenizer_class, **kwargs) - - @classmethod - def _build_tokenizer( - cls, - task: str, - lang: str, - n_model: str, - tokenizer_class: Union[Tokenizer, PreTrainedTokenizerBase], - revision: Optional[str] = None, - cache_dir: Optional[str] = None, - force_download: bool = False, - resume_download: bool = False, - **kwargs, - ) -> Union[Tokenizer, PreTrainedTokenizerBase]: - - if issubclass(tokenizer_class, PreTrainedTokenizerBase): - tokenizer = tokenizer_class.from_pretrained( - pretrained_model_name_or_path=HUB_NAME, - subfolder=f"{task}/{lang}/{n_model}", - revision=revision, - cache_dir=cache_dir, - force_download=force_download, - resume_download=resume_download, - **kwargs, - ) - tokenizer.lang = lang - return tokenizer - - _download_from_hf_hub = partial( - download_from_hf_hub, - model_id=HUB_NAME, - subfolder=f"{task}/{lang}/{n_model}", - revision=revision, - cache_dir=cache_dir, - force_download=force_download, - resume_download=resume_download, - user_agent=TOKENIZER_USER_AGENT, - ) - - # Load from URL or cache if already cached - resolved_vocab_file = _download_from_hf_hub(filename=VOCAB_NAME) - - # _dict_from_json_file - with open(resolved_vocab_file, "r", encoding="utf-8") as f: - vocab = json.load(f) - - if "pos" in tokenizer_class.__name__.lower(): - resolved_pos_vocab_file = _download_from_hf_hub(filename="pos_vocab.json") - - with open(resolved_pos_vocab_file, "r", encoding="utf-8") as f: - pos_vocab = json.load(f) - - kwargs.update({"pos_vocab": pos_vocab}) - - tokenizer = tokenizer_class(lang, vocab, **kwargs) - - if "bpe" in tokenizer_class.__name__.lower(): - encoder = None - bpe_merges = None - - if lang == "en": - encoder_json = _download_from_hf_hub(filename="encoder.json") - with open(encoder_json, "r") as f: - encoder = json.load(f) - - vocab_bpe = _download_from_hf_hub(filename="vocab.bpe") - with open(vocab_bpe, "r", encoding="utf-8") as f: - bpe_data = f.read() - bpe_merges = [ - tuple(merge_str.split()) for merge_str in bpe_data.split("\n")[1:-1] - ] - - tokenizer._build_bpe(lang, encoder, bpe_merges) - - return tokenizer +from .load import DoolyTokenizerHub # noqa +from .load import load_dooly_tokenizer # noqa diff --git a/dooly/tokenizers/base.py b/dooly/tokenizers/base.py index ddb7fba..31e98f8 100644 --- a/dooly/tokenizers/base.py +++ b/dooly/tokenizers/base.py @@ -1,414 +1,17 @@ -import re -import torch -import unicodedata -from abc import abstractmethod -from typing import List, Union, Dict, Set, Optional +import abc +import json +import collections +from typing import Union, List, Optional -from .import_utils import ( +from transformers import PreTrainedTokenizer +from ..utils.import_utils import ( is_available_kss, is_available_nltk, ) -SPACE_NORMALIZER = re.compile(r"\s+") - InputTexts = Union[str, List[str]] TokenizedOutput = Union[List[str], List[List[str]]] -EncodedOutput = Union[List[int], List[List[int]], torch.Tensor] -PaddedOutput = Union[List[List[int]], torch.Tensor] -DecodedOutput = Union[str, List[str]] - - -class _BaseTokenizer: - def __init__( - self, - lang: str, - vocab: Dict[str, int], - cls_token: str = "", - sep_token: str = "", - pad_token: str = "", - unk_token: str = "", - padding_side: str = "right", - max_seq_length: int = 512, - ): - assert padding_side in ["right", "left"] - self.lang = lang - self.vocab = vocab - self.pos_vocab = None - self.id2token = {i: tok for tok, i in vocab.items()} - self.cls_token = cls_token - self.sep_token = sep_token - self.pad_token = pad_token - self.unk_token = unk_token - self.padding_side = padding_side - self.max_seq_length = max_seq_length - - self._langtok_style = "basic" - self.sub_tokenizer = {} - - @property - def cls_token_id(self) -> int: - return self.vocab[self.cls_token] - - @property - def sep_token_id(self) -> int: - return self.vocab[self.sep_token] - - @property - def pad_token_id(self) -> int: - return self.vocab[self.pad_token] - - @property - def unk_token_id(self) -> int: - return self.vocab[self.unk_token] - - @property - def nspecial(self) -> int: - return 4 # cls, sep, pad, unk - - @property - def langtok_style(self): - return self._langtok_style - - @langtok_style.setter - def langtok_style(self, val: str): - self._langtok_style = val - - def _langtok(self, lang: str): - # https://github.com/pytorch/fairseq/blob/master/fairseq/data/multilingual/multilingual_utils.py#L34 - langtok = "" - if self.langtok_style == "basic": - langtok = f"[{lang.upper()}]" - elif self.langtok_style == "mbart": - mapping = {"en": "_XX", "ja": "_XX", "ko": "_KR", "zh": "_CN"} - langtok = f"[{lang + mapping[lang]}]" - elif self.langtok_style == "multilingual": - langtok = f"__{lang}__" - return langtok - - def _set_sub_tokenizer(self, lang: str, tokenizer_object): - self.sub_tokenizer[lang] = tokenizer_object - - def __call__( - self, - text: InputTexts, - text_pair: Optional[InputTexts] = None, - src_lang: Optional[InputTexts] = None, - tgt_lang: Optional[InputTexts] = None, - padding: Union[str, bool] = False, - return_tokens: bool = False, - return_tags: bool = True, - return_tensors: Union[str, bool] = False, - return_attention_mask: bool = True, - add_special_tokens: bool = True, - no_separator: bool = False, - ) -> Union[TokenizedOutput, Dict[str, EncodedOutput]]: - return self.encode( - text=text, - text_pair=text_pair, - src_lang=src_lang, - tgt_lang=tgt_lang, - padding=padding, - return_tokens=return_tokens, - return_tags=return_tags, - return_tensors=return_tensors, - return_attention_mask=return_attention_mask, - add_special_tokens=add_special_tokens, - no_separator=no_separator, - ) - - def _normalize(self, text: str) -> str: - """ Unicode normalization and whitespace removal (often needed for context) """ - text = unicodedata.normalize("NFKC", text) - text = self._normalize_space(text) - return text - - @staticmethod - def _normalize_space(text: str) -> str: - return SPACE_NORMALIZER.sub(" ", text).strip() - - @abstractmethod - def _tokenize(self, text: str, *args, **kwargs) -> List[str]: - pass - - def tokenize( - self, - text: str, - text_pair: Optional[str] = None, - src_lang: Optional[str] = None, - tgt_lang: Optional[str] = None, - return_tags: bool = True, - add_special_tokens: bool = False, - no_separator: bool = False, - ) -> List[str]: - """ - If you want to use `src_lang` and `tgt_lang` parameters, plz overrides! - """ - if self.pos_vocab is None: - return_tags = False - - tokenized = self._tokenize(text) - - if return_tags: - tokenized, tags = tokenized - - if add_special_tokens: - tokenized = [self.cls_token] + tokenized + [self.sep_token] - - if return_tags: - tags = [self.cls_token] + tags + [self.sep_token] - - if text_pair is not None: - tokenized += [self.sep_token] if not no_separator else [] - tokenized_pair = self._tokenize(text_pair) - - if return_tags: - tags += [self.sep_token] if no_separator else [] - tokenized_pair, tags_pair = tokenized_pair - tags += tags_pair - - tokenized += tokenized_pair - - if add_special_tokens: - tokenized += [self.sep_token] - if return_tags: - tags += [self.sep_token] - - if return_tags: - return tokenized, tags - - return tokenized - - def encode_line( - self, - tokenized: List[str], - add_special_tokens: bool = False, - use_pos_vocab: bool = False, - ) -> List[int]: - vocab = self.vocab - if use_pos_vocab and self.pos_vocab is not None: - vocab = self.pos_vocab - - encoded = [] - for token in tokenized: - encoded.append(vocab.get(token, self.unk_token_id)) - - if add_special_tokens: - encoded = [self.cls_token_id] + encoded + [self.sep_token_id] - - return encoded - - def encode( - self, - text: InputTexts, - text_pair: Optional[InputTexts] = None, - src_lang: Optional[InputTexts] = None, - tgt_lang: Optional[InputTexts] = None, - padding: Union[str, bool] = False, - return_tokens: bool = False, - return_tags: bool = True, - return_tensors: Union[str, bool] = False, - return_attention_mask: bool = True, - add_special_tokens: bool = True, - no_separator: bool = False, - ) -> Union[TokenizedOutput, Dict[str, EncodedOutput]]: - """ Encode tokens to ids, used for single or batched sentence """ - - assert isinstance(return_tensors, bool) or return_tensors == "pt" - return_tensors = (return_tensors == "pt") or return_tensors - - assert text_pair is None or type(text) == type(text_pair) - - if (src_lang is None) ^ (tgt_lang is None): - src_lang = tgt_lang = None - - if not hasattr(self, "pos_tagger"): - return_tags = False - - if isinstance(text, str): - return self.encode( - text=[text], - text_pair=[text_pair], - src_lang=[src_lang], - tgt_lang=[tgt_lang], - padding=padding, - return_tokens=return_tokens, - return_tags=return_tags, - return_tensors=return_tensors, - return_attention_mask=return_attention_mask, - add_special_tokens=add_special_tokens, - no_separator=no_separator, - ) - - if text_pair is None: - text_pair = [None] * len(text) - if src_lang is None: - src_lang = [None] * len(text) - if tgt_lang is None: - tgt_lang = [None] * len(text) - - assert len(text) == len(text_pair) - assert len(src_lang) == len(tgt_lang) - - if len(src_lang) == 1: - src_lang = src_lang * len(text) - tgt_lang = tgt_lang * len(text) - - assert len(text) == len(src_lang) - - texts, text_pairs = text, text_pair - src_langs, tgt_langs = src_lang, tgt_lang - input_ids = [] - segment_labels = [] - - for text, text_pair, src_lang, tgt_lang in zip( - texts, text_pairs, src_langs, tgt_langs - ): - tokenized = self.tokenize( - text=text, - text_pair=text_pair, - src_lang=src_lang, - tgt_lang=tgt_lang, - return_tags=return_tags, - no_separator=no_separator, - add_special_tokens=add_special_tokens, - ) - encoded = None - encoded_tags = None - - if return_tags: - tokenized, tags = tokenized - - if not return_tokens: - encoded = self.encode_line(tokenized=tokenized) - - if return_tags: - encoded_tags = self.encode_line(tokenized=tags, use_pos_vocab=True) - - input_ids.append(tokenized if return_tokens else encoded) - - if return_tags: - segment_labels.append(tags if return_tokens else encoded_tags) - - if return_tokens: - input_ids = input_ids if len(texts) > 1 else input_ids[0] - - if return_tags: - segment_labels = segment_labels if len(texts) > 1 else segment_labels[0] - return input_ids, segment_labels - - return input_ids - - attention_mask = None - if return_tensors or padding: - padded = self.pad( - sequences={"input_ids": input_ids}, - padding=padding, - return_tensors=return_tensors, - ) - input_ids = padded["input_ids"] - attention_mask = padded["attention_mask"] - - if return_tags: - segment_labels = self.pad( - sequences={"input_ids": segment_labels}, - padding=padding, - return_tensors=return_tensors, - )["input_ids"] - - batch_encoding = {"input_ids": input_ids} - - if return_attention_mask and attention_mask is not None: - batch_encoding.update({"attention_mask": attention_mask}) - - if return_tags: - batch_encoding.update({"segment_labels": segment_labels}) - - return batch_encoding - - def decode_line(self, ids: List[int], ignore_symbols: Set[int] = {}) -> str: - sent = [] - for _id in ids: - if _id not in ignore_symbols: - sent.append(self.id2token.get(_id, self.unk_token)) - return " ".join(sent) - - def _recover_original(self, decoded_text: str) -> str: - return decoded_text - - def decode( - self, - ids: EncodedOutput, - ignore_symbols: List[int] = [], - recover_original: bool = True, - ) -> DecodedOutput: - - if isinstance(ids, torch.Tensor): - ids = ids.detach().cpu().tolist() - - if isinstance(ids[0], int): - return self.decode( - ids=[ids], - ignore_symbols=ignore_symbols, - recover_original=recover_original, - ) - - ignore_symbols = set(None or ignore_symbols) - ignore_symbols.update([self.cls_token_id, self.sep_token_id, self.pad_token_id]) - - list_of_ids = ids - decoded_texts = [] - for ids in list_of_ids: - decoded = self.decode_line(ids, ignore_symbols) - if recover_original: - decoded = self._recover_original(decoded) - decoded_texts.append(decoded) - - if len(decoded_texts) == 1: - decoded_texts = decoded_texts[0] - - return decoded_texts - - def pad( - self, - sequences: Dict[str, EncodedOutput], - padding: Union[str, bool] = True, - return_tensors: bool = True, - pad_to_multiple_of: Union[int, bool] = False, # match to hf pad method - ) -> Dict[str, PaddedOutput]: - """Pad batched sequences. - if return_tensors, then return torch.LongTensor object. - """ - - input_ids = sequences.get("input_ids") - assert input_ids is not None - - if isinstance(input_ids[0], int): - input_ids = [input_ids] - - max_length = -1 - if padding == "max_length": - max_length = self.max_seq_length - else: - max_length = max(len(ids) for ids in input_ids) - - padded = {"input_ids": [], "attention_mask": []} - for ids in input_ids: - seq_len = len(ids) - if self.padding_side == "right": - ids = ids + [self.pad_token_id] * (max_length - seq_len) - attn_mask = [1] * seq_len + [0] * (max_length - seq_len) - else: - ids = [self.pad_token_id] * (max_length - seq_len) + ids - attn_mask = [0] * (max_length - seq_len) + [1] * seq_len - padded["input_ids"].append(ids) - padded["attention_mask"].append(attn_mask) - - if return_tensors: - for k, v in padded.items(): - padded[k] = torch.LongTensor(v) - - return padded class SentTokenizeMixin: @@ -440,9 +43,7 @@ def _set_sent_tokenizer(self): ) def sent_tokenize( - self, - texts: InputTexts, - langs: Optional[InputTexts] = None, + self, texts: InputTexts, langs: Optional[InputTexts] = None ) -> List[List[str]]: if isinstance(texts, str): texts = [texts] @@ -492,7 +93,64 @@ def sent_tokenize( return sentences, num_sentences -class Tokenizer(_BaseTokenizer, SentTokenizeMixin): - """ Whitespace Base Tokenizer with sentence tokenizer """ +class DoolyPreTrainedTokenizer(PreTrainedTokenizer, SentTokenizeMixin): + vocab_files_names = {"vocab_file": "vocab.json"} + replacement: Optional[str] = None + + def __init__( + self, + vocab_file, + cls_token: str = "", + sep_token: str = "", + pad_token: str = "", + unk_token: str = "", + task: str = None, + lang: str = None, + n_model: str = None, + **kwargs + ): + super().__init__( + cls_token=cls_token, + sep_token=sep_token, + pad_token=pad_token, + unk_token=unk_token, + **kwargs, + ) + with open(vocab_file, "r", encoding="utf-8") as f: + self.vocab = json.load(f) + self.ids_to_tokens = collections.OrderedDict( + [(ids, tok) for tok, ids in self.vocab.items()] + ) + + replacement = kwargs.pop("replacement", None) + self.replacement = replacement or self.replacement + + self.task = task + self.lang = lang + self.n_model = n_model + + @abc.abstractmethod + def _tokenize(self, text: str, **kwargs) -> List[str]: + pass + + @property + def vocab_size(self) -> int: + return len(self.vocab) + + def get_vocab(self): + return dict(self.vocab, **self.added_tokens_encoder) + + def build_inputs_with_special_tokens( + self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None + ) -> List[int]: + if token_ids_1 is None: + return [self.cls_token_id] + token_ids_0 + [self.sep_token_id] + cls = [self.cls_token_id] + sep = [self.sep_token_id] + return cls + token_ids_0 + sep + token_ids_1 + sep + + def _convert_token_to_id(self, token: str) -> int: + return self.vocab.get(token, self.vocab.get(self.unk_token)) - pass + def _convert_id_to_token(self, index: int) -> str: + return self.ids_to_tokens.get(index, self.unk_token) diff --git a/dooly/tokenizers/bpe.py b/dooly/tokenizers/bpe.py index e9abd5a..8fc56c1 100644 --- a/dooly/tokenizers/bpe.py +++ b/dooly/tokenizers/bpe.py @@ -1,14 +1,17 @@ -from typing import Dict, List, Set, Union +from typing import List from functools import lru_cache -import torch +import transformers +from transformers import GPT2TokenizerFast +from transformers import BertTokenizer, BertJapaneseTokenizer -from .base import Tokenizer - - -TokenizedOutput = Union[List[str], List[List[str]]] -EncodedOutput = Union[List[int], List[List[int]], torch.Tensor] -DecodedOutput = Union[str, List[str]] +from .load import load_dooly_tokenizer +from .base import DoolyPreTrainedTokenizer +from ..utils import ( + recover_original_hf_bucket_url, + is_available_ipadic, + is_available_fugashi, +) """ @@ -45,188 +48,103 @@ def bytes_to_unicode(): return dict(zip(bs, cs)) -def get_pairs(word): - """Return set of symbol pairs in a word. - Word is represented as tuple of symbols (symbols being variable-length strings). - """ - pairs = set() - prev_char = word[0] - for char in word[1:]: - pairs.add((prev_char, char)) - prev_char = char - return pairs - - -class Encoder: - def __init__(self, encoder, bpe_merges, errors="replace"): - self.encoder = encoder - self.decoder = {v: k for k, v in self.encoder.items()} - self.errors = errors # how to handle errors in decoding - self.byte_encoder = bytes_to_unicode() - self.byte_decoder = {v: k for k, v in self.byte_encoder.items()} - self.bpe_ranks = dict(zip(bpe_merges, range(len(bpe_merges)))) - self.cache = {} - - try: - import regex as re - - self.re = re - except ImportError: - raise ImportError("Please install regex with: pip install regex") - - # Should haved added re.IGNORECASE so BPE merges can happen for capitalized versions of contractions - self.pat = self.re.compile( - r"""'s|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+""" - ) +class DoolyGPT2TokenizerFast(DoolyPreTrainedTokenizer): + vocab_files_names = { + "vocab_file": "vocab.json", + "bpe_merge_file": "vocab.bpe", + "bpe_vocab_file": "encoder.json", + } - def bpe(self, token): - if token in self.cache: - return self.cache[token] - word = tuple(token) - pairs = get_pairs(word) - - if not pairs: - return token - - while True: - bigram = min(pairs, key=lambda pair: self.bpe_ranks.get(pair, float("inf"))) - if bigram not in self.bpe_ranks: - break - first, second = bigram - new_word = [] - i = 0 - while i < len(word): - try: - j = word.index(first, i) - new_word.extend(word[i:j]) - i = j - except: # noqa - new_word.extend(word[i:]) - break - - if word[i] == first and i < len(word) - 1 and word[i + 1] == second: - new_word.append(first + second) - i += 2 - else: - new_word.append(word[i]) - i += 1 - new_word = tuple(new_word) - word = new_word - if len(word) == 1: - break - else: - pairs = get_pairs(word) - word = " ".join(word) - self.cache[token] = word - return word - - def encode(self, text: str) -> List[str]: - bpe_tokens = [] - for token in self.re.findall(self.pat, text): - token = "".join(self.byte_encoder[b] for b in token.encode("utf-8")) - # bpe_tokens.extend( - # self.encoder[bpe_token] for bpe_token in self.bpe(token).split(" ") - # ) - bpe_tokens.extend(self.bpe(token).split(" ")) - return bpe_tokens - - def decode(self, tokens): - text = "".join([self.decoder.get(token, token) for token in tokens]) - text = self._decode(text) - return text - - def _decode(self, text: str) -> str: - return bytearray([self.byte_decoder[c] for c in text]).decode( + def __init__( + self, + bpe_vocab_file=None, + bpe_merge_file=None, + errors: str = "replace", + **kwargs + ): + super().__init__(**kwargs) + self.bpe = GPT2TokenizerFast( + vocab_file=bpe_vocab_file, merges_file=bpe_merge_file + ) + self.bpe_vocab = self.bpe.backend_tokenizer.get_vocab() + self.byte_decoder = {v: k for k, v in bytes_to_unicode().items()} + self.errors = errors + + @property + def bpe_vocab_size(self) -> int: + return len(self.get_bpe_vocab()) + + def get_bpe_vocab(self): + return self.bpe_vocab + + def _tokenize(self, text: str, **kwargs) -> List[str]: + return self.bpe.tokenize(text) + + def _convert_token_to_id(self, token: str) -> int: + bpe_token = token + if token not in self.all_special_tokens: + bpe_token = self.bpe._tokenizer.token_to_id(token) + return self.vocab.get(str(bpe_token), self.vocab.get(self.unk_token)) + + def _convert_id_to_token(self, index: int) -> str: + bpe_id = self.ids_to_tokens.get(index, self.unk_token) + if bpe_id not in self.all_special_tokens: + bpe_id = int(bpe_id) + return self.bpe._tokenizer.id_to_token(bpe_id) + + def convert_tokens_to_string(self, tokens: List[str]) -> str: + """Converts a sequence of tokens (string) in a single string.""" + text = "".join(tokens) + text = bytearray([self.byte_decoder[c] for c in text]).decode( "utf-8", errors=self.errors ) + return text -class Gpt2BpeTokenizer(Tokenizer): - """ GPT2 BytePairEncoding Tokenizer """ - - def _build_bpe(self, lang: str, encoder_json: Dict = None, bpe_merges: Dict = None): - self._bpe = Encoder(encoder_json, bpe_merges) +class DoolyBertTokenizer(DoolyPreTrainedTokenizer): + replacement: str = "##" + wp_path: str = None + wp_tok_class: transformers.PreTrainedTokenizer = None - def __call__( - self, *args, **kwargs - ) -> Union[TokenizedOutput, Dict[str, EncodedOutput]]: - return_tokens = kwargs.pop("return_tokens", False) - add_special_tokens = kwargs.pop("add_special_tokens", False) + def __init__(self, **kwargs): + super().__init__(**kwargs) + with recover_original_hf_bucket_url(): + wp_path = kwargs.pop("wp_path", self.wp_path) + self.wordpiece = load_dooly_tokenizer(wp_path, self.wp_tok_class) - if return_tokens and add_special_tokens: - add_special_tokens = False + def _tokenize(self, text: str, **kwargs) -> List[str]: + return self.wordpiece.tokenize(text) - kwargs.update( - {"return_tokens": return_tokens, "add_special_tokens": add_special_tokens} + def convert_tokens_to_string(self, tokens: List[str]) -> str: + """Converts a sequence of tokens (string) in a single string.""" + text = "".join( + [token.replace(" ", "").replace(self.replacement, "") for token in tokens] ) + return text - outputs = self.encode(*args, **kwargs) - - if return_tokens: - _outputs = [] - if not isinstance(outputs[0], list): - outputs = [outputs] - for output in outputs: - _outputs.append([self._bpe._decode(o) for o in output]) - if len(_outputs) == 1: - _outputs = _outputs[0] - outputs = _outputs - - return outputs - def _tokenize(self, text: str) -> List[str]: - # return list(map(str, self._bpe.encode(text))) - return self._bpe.encode(text) - - def encode_line( - self, - tokenized: List[str], - add_special_tokens: bool = False, - use_pos_vocab: bool = False, - ) -> List[int]: - encoded = [] - for bpe_token in tokenized: - if bpe_token not in ["", "", "", ""]: - bpe_token = str(self._bpe.encoder[bpe_token]) - encoded.append(self.vocab.get(bpe_token, self.unk_token_id)) - - if add_special_tokens: - encoded = [self.cls_token_id] + encoded + [self.sep_token_id] - return encoded - - def decode_line(self, ids: List[int], ignore_symbols: Set[int] = {}) -> str: - x = super().decode_line(ids, ignore_symbols) - return self._bpe.decode( - [int(tok) if tok not in {"", ""} else tok for tok in x.split()] - ) +class DoolyBertJaTokenizer(DoolyBertTokenizer): + wp_path: str = "cl-tohoku/bert-base-japanese-whole-word-masking" + wp_tok_class: transformers.PreTrainedTokenizer = BertJapaneseTokenizer + def __init__(self, **kwargs): + if is_available_ipadic(): + import ipadic # noqa + else: + raise ModuleNotFoundError( + "Please install ipadic with: `pip install ipadic`" + ) -class BpeJaZhTokenizer(Tokenizer): - """ BytePairEncoding Tokenizer for Ja and Zh """ - - def _build_bpe(self, lang: str, encoder_json: Dict = None, bpe_merges: Dict = None): - if lang == "ja": - try: - import ipadic # noqa - except ImportError: - raise ImportError("Please install ipadic with: `pip install ipadic`") - try: - import fugashi # noqa - except ImportError: - raise ImportError("Please install fugashi with: `pip install fugashi`") - from transformers import BertJapaneseTokenizer - - model_name_or_path = "cl-tohoku/bert-base-japanese-whole-word-masking" - self._bpe = BertJapaneseTokenizer.from_pretrained(model_name_or_path) - elif lang == "zh": - from transformers import BertTokenizer - - self._bpe = BertTokenizer.from_pretrained( - "bert-base-chinese", do_lower_case=True + if is_available_fugashi(): + import fugashi # noqa + else: + raise ModuleNotFoundError( + "Please install fugashi with: `pip install fugashi`" ) - def _tokenize(self, text: str) -> List[str]: - return self._bpe.tokenize(text) + super().__init__(**kwargs) + - def _recover_original(self, decoded_text: str) -> str: - return decoded_text.replace(" ", "") +class DoolyBertZhTokenizer(DoolyBertTokenizer): + wp_path: str = "bert-base-chinese" + wp_tok_class: transformers.PreTrainedTokenizer = BertTokenizer diff --git a/dooly/tokenizers/char.py b/dooly/tokenizers/char.py index 9809284..7fc0f2f 100644 --- a/dooly/tokenizers/char.py +++ b/dooly/tokenizers/char.py @@ -1,81 +1,257 @@ -from typing import List, Optional +import re +from typing import Union, List, Dict, Optional -from .base import Tokenizer +import torch +from transformers import BatchEncoding, TensorType +from .base import DoolyPreTrainedTokenizer -class CharS1Tokenizer(Tokenizer): - """ Character Tokenizer with style 1 """ - def _recover_original(self, decoded_text: str) -> str: - return decoded_text.replace(" ", "").replace("▁", " ").strip() +class DoolyCharTokenizer(DoolyPreTrainedTokenizer): + replacement: str = "▁" + __SPACE_NORMALIZER__ = re.compile(r"\s+") - def _tokenize(self, text: str): + def _normalize_space(self, text: str) -> str: + return self.__SPACE_NORMALIZER__.sub(" ", text).strip() + + def convert_tokens_to_string(self, tokens: List[str]) -> str: + """Converts a sequence of tokens (string) in a single string.""" + text = "".join(tokens) + return text.replace(" ", "").replace(self.replacement, " ").strip() + + def _tokenize_chatbpe_style(self, text: str) -> List[str]: + """e.g., + >>> text = "손흥민은 28세의 183 센티미터, 77 킬로그램이며, 현재 주급은 약 3억 원이다." + >>> tokenizer._tokenize_charbpe_style(text) + ['▁손', '흥', '민', '은', '▁2', '8', '세', '의', '▁1', '8', '3', '▁센', '티', '미', '터', + ',', '▁7', '7', '▁킬', '로', '그', '램', '이', '며', ',', '▁현', '재', '▁주', '급', '은', + '▁약', '▁3', '억', '▁원', '이', '다', '.'] + """ x = text.strip() x = [c for c in self._normalize_space(x)] tokenized = list() for i in range(len(x)): if x[i] == " ": - x[i + 1] = f"▁{x[i+1]}" + x[i + 1] = self.replacement + f"{x[i + 1]}" continue else: tokenized.append(x[i]) - tokenized[0] = f"▁{tokenized[0]}" + tokenized[0] = self.replacement + f"{tokenized[0]}" return tokenized - -class CharS2Tokenizer(Tokenizer): - """ Character Tokenizer with style 2 """ - - def _recover_original(self, decoded_text: str) -> str: - return decoded_text.replace(" ", "").replace("▁", " ").strip() - - def _tokenize(self, text: str) -> List[str]: + def _tokenize_whitespace_style(self, text: str) -> List[str]: + """e.g., + >>> text = "손흥민은 28세의 183 센티미터, 77 킬로그램이며, 현재 주급은 약 3억 원이다." + >>> tokenizer._tokenize_whitespace_style(text) + ['손', '흥', '민', '은', '▁', '2', '8', '세', '의', '▁', '1', '8', '3', '▁', '센', '티', + '미', '터', ',', '▁', '7', '7', '▁', '킬', '로', '그', '램', '이', '며', ',', '▁', '현', + '재', '▁', '주', '급', '은', '▁', '약', '▁', '3', '억', '▁', '원', '이', '다', '.'] + """ text = text.strip() - text = text.replace(" ", "▁") + text = text.replace(" ", self.replacement) text = " ".join([c for c in text]) - tokenized = self._normalize_space(text) - tokenized = tokenized.split() + text = self._normalize_space(text) + tokenized = text.split() return tokenized - def tokenize( + +class DoolyCharBertTokenizer(DoolyCharTokenizer): + def _tokenize(self, text: str, **kwargs) -> List[str]: + return self._tokenize_chatbpe_style(text) + + +# TODO: WSD를 위한 target text tokenize function 작성 +# PORORO WSD transformer: +# src_tokens: token_ids_0 + [] +# tgt_tokens: ??? -> 분석 필요 +class DoolyCharSeq2SeqWsdTokenizer(DoolyCharTokenizer): + def _tokenize(self, text: str, **kwargs) -> List[str]: + return self._tokenize_whitespace_style(text) + + def build_inputs_with_special_tokens( + self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None + ) -> List[int]: + sep = [self.sep_token_id] + if token_ids_1 is None: + return token_ids_0 + sep + return token_ids_0 + sep + token_ids_1 + sep + + +class DoolyCharSeq2SeqNmtTokenizer(DoolyCharTokenizer): + """ + - mBART: + src_tokens: token_ids_0 + [] + [src_lang_code] + tgt_tokens: token_ids_1 + [] + [tgt_lang_code] + - PORORO NMT transformer: + src_tokens: [src_lang_code] + token_ids_0 + [tgt_lang_code] + [] + tgt_tokens: [] + token_ids_1 + [] + """ + + __LANG_TO_CODE__ = { + "ko": "[ko_KR]", + "en": "[en_XX]", + "ja": "[ja_XX]", + "zh": "[zh_CN]", + } + + def __init__(self, **kwargs): + super().__init__( + additional_special_tokens=list(self.lang_to_code.values()), **kwargs + ) + + def _tokenize(self, text: str, **kwargs) -> List[str]: + return self._tokenize_whitespace_style(text) + + def __call__( self, - text: str, - text_pair: Optional[str] = None, - src_lang: Optional[str] = None, - tgt_lang: Optional[str] = None, - add_special_tokens: bool = True, - no_separator: bool = False, + text, + text_pair, + src_lang: Union[str, List[str]] = None, + tgt_lang: Union[str, List[str]] = None, **kwargs, - ) -> List[str]: - if (src_lang is None) ^ (tgt_lang is None): - src_lang = tgt_lang = None + ) -> BatchEncoding: + if src_lang is None and tgt_lang is None: + return super().__call__(text, text_pair, **kwargs) - text = text.strip() - if self.sub_tokenizer.get(src_lang, None) is not None: - sub_tokenizer = self.sub_tokenizer[src_lang] - if hasattr(sub_tokenizer, "segment"): - tokenized = sub_tokenizer.segment(text) - elif hasattr(sub_tokenizer, "tokenize"): - tokenized = sub_tokenizer.tokenize(text, add_special_tokens=False) - else: - raise AttributeError - else: - tokenized = self._tokenize(text) + assert text_pair is not None - if src_lang is not None: - tokenized = [self._langtok(src_lang)] + tokenized - if tgt_lang is not None: - tokenized = tokenized + [self._langtok(tgt_lang)] + add_special_tokens = kwargs.pop("add_special_tokens", True) + return_tensors = kwargs.get("return_tensors", None) - if add_special_tokens: - tokenized += [self.sep_token] + # tokenize source text + batch_encodings = super().__call__(text, add_special_tokens=False, **kwargs) + batch_encodings = self.add_language_tokens( + batch_encodings, src_lang, tgt_lang, add_special_tokens, return_tensors + ) - if text_pair is not None: - tokenized += self._tokenize(text_pair) + # tokenize target text + label_encodings = super().__call__( + text_pair, add_special_tokens=add_special_tokens, **kwargs + ) - if add_special_tokens: - tokenized += [self.sep_token] + batch_encodings["labels"] = label_encodings["input_ids"] + return batch_encodings - return tokenized + @property + def lang_to_code(self) -> Dict[str, str]: + return self.__LANG_TO_CODE__ + + @property + def lang_code_to_id(self) -> Dict[str, int]: + return { + code: self.vocab.get(code, self.unk_token_id) + for code in self.lang_to_code.values() + } + + def lang_to_id(self, lang: str) -> int: + lang_code = self.lang_to_code.get(lang, self.unk_token) + return self.lang_code_to_id.get(lang_code, self.unk_token_id) + + def add_language_tokens( + self, + batch_encodings: BatchEncoding, + src_lang: Union[str, List[str]], + tgt_lang: Union[str, List[str]], + add_special_tokens: bool = False, + return_tensors: Optional[Union[str, TensorType]] = None, + ) -> BatchEncoding: + if return_tensors is not None and return_tensors != "pt": + raise ValueError("Only support tensor type `pt`.") + + input_ids = batch_encodings.pop("input_ids") + attention_mask = batch_encodings.get("attention_mask", None) + token_type_ids = batch_encodings.get("token_type_ids", None) + + src_langs = src_lang + if isinstance(src_lang, str): + src_langs = [src_lang] * len(input_ids) + tgt_langs = tgt_lang + if isinstance(tgt_lang, str): + tgt_langs = [tgt_lang] * len(input_ids) + + assert len(src_langs) == len(input_ids) + assert len(tgt_langs) == len(input_ids) + + token_added_ids = [] + if attention_mask is not None: + token_added_masks = [] + if token_type_ids is not None: + token_added_type_ids = [] + + for i in range(len(input_ids)): + _input_ids = input_ids[i] + if attention_mask is not None: + _attention_mask = attention_mask[i] + if token_type_ids is not None: + _token_type_ids = token_type_ids[i] + src_lang = src_langs[i] + tgt_lang = tgt_langs[i] + + maximum_idx = [ + i for i, val in enumerate(_input_ids) if val != self.pad_token_id + ] + idx_to_add = 0 + if len(maximum_idx) > 0: + idx_to_add = max(maximum_idx) + 1 + + src_lang = self.lang_to_id(src_lang) + tgt_lang = self.lang_to_id(tgt_lang) + sep = self.sep_token_id + + _input_ids = self.insert_tokens( + _input_ids, [src_lang], [tgt_lang, sep], idx_to_add, return_tensors + ) + token_added_ids.append(_input_ids) + + if attention_mask is not None: + _attention_mask = self.insert_tokens( + _attention_mask, [1], [1, 1], idx_to_add, return_tensors + ) + token_added_masks.append(_attention_mask) + + if token_type_ids is not None: + _token_type_ids = self.insert_tokens( + _token_type_ids, [0], [0, 0], idx_to_add, return_tensors + ) + token_added_type_ids.append(_token_type_ids) + + def unsqueeze_and_cat(tensorlist: List[torch.Tensor]) -> torch.Tensor: + dims = [0] * len(tensorlist) + tensorlist = list(map(torch.unsqueeze, tensorlist, dims)) + return torch.cat(tensorlist) + + if return_tensors: + token_added_ids = unsqueeze_and_cat(token_added_ids) + if attention_mask is not None: + token_added_masks = unsqueeze_and_cat(token_added_masks) + if token_type_ids is not None: + token_added_type_ids = unsqueeze_and_cat(token_added_type_ids) + + batch_encodings["input_ids"] = token_added_ids + if attention_mask is not None: + batch_encodings["attention_mask"] = token_added_masks + if token_type_ids is not None: + batch_encodings["token_type_ids"] = token_added_type_ids + return batch_encodings + + def insert_tokens( + self, + ids: Union[List[int], torch.Tensor], + prefix: List[int], + suffix: List[int], + idx_to_add: int, + return_tensors: Optional[Union[str, TensorType]] = None, + ) -> Union[List[int], torch.Tensor]: + if return_tensors: + ids = torch.cat( + [ + torch.tensor(prefix, requires_grad=False), + ids[:idx_to_add], + torch.tensor(suffix, requires_grad=False), + ids[idx_to_add:], + ] + ).long() + return ids + return prefix + ids[:idx_to_add] + suffix + ids[idx_to_add:] diff --git a/dooly/tokenizers/hf_tokenizer.py b/dooly/tokenizers/fast.py similarity index 79% rename from dooly/tokenizers/hf_tokenizer.py rename to dooly/tokenizers/fast.py index f62221f..eb2303e 100644 --- a/dooly/tokenizers/hf_tokenizer.py +++ b/dooly/tokenizers/fast.py @@ -1,6 +1,5 @@ from typing import List, Optional from tokenizers import Encoding -from transformers import PreTrainedTokenizerBase # noqa from transformers import RobertaTokenizerFast as _RobertaTokenizerFast from transformers import PreTrainedTokenizerFast as _PreTrainedTokenizerFast @@ -39,26 +38,19 @@ def build_custom_roberta_tokenizer( import tokenizers bpe_obj = tokenizers.models.BPE.from_file( - vocab_filename, - merges_filename, - unk_token="", - fuse_unk=True, + vocab_filename, merges_filename, unk_token="", fuse_unk=True ) # @TODO: Unigram _tokenizer = tokenizers.Tokenizer(bpe_obj) _tokenizer.normalizer = tokenizers.normalizers.NFKC() _tokenizer.pre_tokenizer = tokenizers.pre_tokenizers.Metaspace( - replacement=replacement, - add_prefix_space=add_prefix_space, + replacement=replacement, add_prefix_space=add_prefix_space ) _tokenizer.post_processor = tokenizers.processors.RobertaProcessing( - sep=("", 2), - cls=("", 0), - add_prefix_space=False, + sep=("", 2), cls=("", 0), add_prefix_space=False ) _tokenizer.decoder = tokenizers.decoders.Metaspace( - replacement=replacement, - add_prefix_space=add_prefix_space, + replacement=replacement, add_prefix_space=add_prefix_space ) return RobertaTokenizerFast( @@ -68,18 +60,11 @@ def build_custom_roberta_tokenizer( ) -# To match the class name to avoid warning statements -# when `config_tokenizer_class` is not None. -# See here: transformers 133c5e40 -# ./src/transformers/tokenization_utils_base.py#L1825 class RobertaTokenizerFast(_RobertaTokenizerFast, SentTokenizeMixin): def segment(self, texts: InputTexts) -> TokenizedOutput: if isinstance(texts, str): texts = [texts] - encodings = self.backend_tokenizer.encode_batch( - texts, - add_special_tokens=False, - ) + encodings = self.backend_tokenizer.encode_batch(texts, add_special_tokens=False) results = [] for text, encoding in zip(texts, encodings): results.append(self._unk_to_raw_text(text, encoding)) diff --git a/dooly/tokenizers/import_utils.py b/dooly/tokenizers/import_utils.py deleted file mode 100644 index 2c79f9c..0000000 --- a/dooly/tokenizers/import_utils.py +++ /dev/null @@ -1,31 +0,0 @@ -import os -import importlib - - -def is_available_mecab(): - _mecab = importlib.util.find_spec("mecab") - if os.name != "nt": - return _mecab is not None - else: - _eunjeon = importlib.util.find_spec("eunjeon") - return _mecab is not None and _eunjeon is not None - - -def is_available_ipadic(): - return importlib.util.find_spec("ipadic") - - -def is_available_fugashi(): - return importlib.util.find_spec("fugashi") - - -def is_available_jieba(): - return importlib.util.find_spec("jieba") - - -def is_available_nltk(): - return importlib.util.find_spec("nltk") - - -def is_available_kss(): - return importlib.util.find_spec("kss") diff --git a/dooly/tokenizers/load.py b/dooly/tokenizers/load.py new file mode 100644 index 0000000..d05bd8a --- /dev/null +++ b/dooly/tokenizers/load.py @@ -0,0 +1,116 @@ +from typing import Type, Union + +import transformers +from ..utils import _locate, DOOLY_HUB_NAME + + +DoolyTokenizerHub = { + "dp": {"ko": {"posbert.base": "pos_tagger.DoolyPosDpTokenizer"}}, + "mrc": {"ko": {"brainbert.base": "fast.RobertaTokenizerFast"}}, + "mt": { + "multi": { + "transformer.large.mtpg": "char.DoolyCharSeq2SeqNmtTokenizer", + "transformer.large.fast.mtpg": "char.DoolyCharSeq2SeqNmtTokenizer", + }, + }, + "ner": { + "ko": {"charbert.base": "char.DoolyCharBertTokenizer"}, + "en": {"roberta.base": "bpe.DoolyGPT2TokenizerFast"}, + "ja": {"jaberta.base": "bpe.DoolyBertJaTokenizer"}, + "zh": {"zhberta.base": "bpe.DoolyBertZhTokenizer"}, + }, + "nli": { + "ko": {"brainbert.base": "fast.RobertaTokenizerFast"}, + "en": {"roberta.base": "bpe.DoolyGPT2TokenizerFast"}, + "ja": {"jaberta.base": "bpe.DoolyBertJaTokenizer"}, + "zh": {"zhberta.base": "bpe.DoolyBertZhTokenizer"}, + }, + "qg": {"ko": {"kobart.base": "bpe.PreTrainedTokenizerFast"}}, + "wsd": {"ko": {"transformer.large": "char.DoolyCharSeq2SeqWsdTokenizer"}}, +} +DoolyTokenizerHub["bt"] = DoolyTokenizerHub["mt"] +DoolyTokenizerHub["zero_topic"] = DoolyTokenizerHub["nli"] + +available_tasks = list(DoolyTokenizerHub.keys()) + + +def load_pretrained_tokenizer( + pretrained_model_name_or_path: str, + tokenizer_class: Type[transformers.PreTrainedTokenizer], + **kwargs, +) -> transformers.PreTrainedTokenizer: + return tokenizer_class.from_pretrained(pretrained_model_name_or_path, **kwargs) + + +def load_tokenizer_from_dooly_hub( + subfolder: str, tokenizer_class: Type[transformers.PreTrainedTokenizer], **kwargs +) -> transformers.PreTrainedTokenizer: + def _load_pretrained(pretrained_model_name_or_path: str, subfolder: str, **kwargs): + return tokenizer_class.from_pretrained( + pretrained_model_name_or_path, subfolder=subfolder, **kwargs + ) + + return _load_pretrained( + pretrained_model_name_or_path=DOOLY_HUB_NAME, subfolder=subfolder, **kwargs + ) + + +def load_dooly_tokenizer( + pretrained_model_name_or_path: str = None, + tokenizer_class: Union[str, Type[transformers.PreTrainedTokenizer]] = None, + task: str = None, + lang: str = None, + n_model: str = None, + **kwargs, +) -> transformers.PreTrainedTokenizer: + if pretrained_model_name_or_path is not None: + if tokenizer_class is None: + raise ValueError( + "If you are using the personal huggingface.co model, " + "`tokenizer_class` parameter is required." + ) + return load_pretrained_tokenizer( + pretrained_model_name_or_path, tokenizer_class, **kwargs + ) + + if all([task is None and lang is None and n_model is None]): + raise ValueError( + "`task`, `lang`, and `n_model` parameters are required to " + "access the subfolder of dooly-hub.\nCheck your parameters! " + f"`task`: {task} `lang`: {lang} `n_model`: {n_model}." + ) + + assert ( + task in available_tasks + ), f"Task `{task}` is not available. See here {available_tasks}." + + available_langs = DoolyTokenizerHub[task] + assert lang in available_langs, ( + f"Language `{lang}` is not available in this task {task}. " + f"See here {available_langs}." + ) + + available_models = available_langs[lang] + assert n_model in available_models, ( + f"Model `{n_model}` is not available in this task-lang pair. " + f"See here {available_models}." + ) + + subfolder = f"{task}/{lang}/{n_model}" + subfolder_postfix = kwargs.pop("subfolder_postfix", None) + if subfolder_postfix is not None: + subfolder += f"/{subfolder_postfix}" + + if tokenizer_class is None: + module_path = "dooly.tokenizers." + available_models[n_model] + elif isinstance(tokenizer_class, str): + module_path = tokenizer_class + + if not issubclass(tokenizer_class, transformers.PreTrainedTokenizer): + tokenizer_class = _locate(module_path) + + kwargs.update({"task": task, "lang": lang, "n_model": n_model}) + + return load_tokenizer_from_dooly_hub( + subfolder=subfolder, tokenizer_class=tokenizer_class, **kwargs + ) diff --git a/dooly/tokenizers/pos_tagger.py b/dooly/tokenizers/pos_tagger.py index e4e420e..847e105 100644 --- a/dooly/tokenizers/pos_tagger.py +++ b/dooly/tokenizers/pos_tagger.py @@ -1,12 +1,14 @@ -import re import os +import re import abc -from typing import List, Tuple, Union, Dict +import json +import inspect +from typing import List, Tuple, Union, Dict, Callable, Any -import torch +from transformers import BatchEncoding -from .base import Tokenizer -from .import_utils import ( +from .base import DoolyPreTrainedTokenizer +from ..utils.import_utils import ( is_available_mecab, is_available_ipadic, is_available_fugashi, @@ -15,11 +17,6 @@ ) -InputTexts = Union[str, List[str]] -TokenizedOutput = Union[List[str], List[List[str]]] -EncodedOutput = Union[List[int], List[List[int]], torch.Tensor] -PaddedOutput = Union[List[List[int]], torch.Tensor] - PosTagResult = Union[Tuple[str, str], str] @@ -249,25 +246,29 @@ def pos(self, sent: str, **kwargs) -> PosTagResult: } -class PosDpTokenizer(Tokenizer): - def __init__(self, *args, **kwargs): - pos_vocab: Dict[str, int] = kwargs.pop("pos_vocab", None) +class DoolyPosDpTokenizer(DoolyPreTrainedTokenizer): + vocab_files_names = { + "vocab_file": "vocab.json", + "pos_vocab_file": "pos_vocab.json", + } + replacement: str = "▃" - super().__init__(*args, **kwargs) + def __init__(self, pos_vocab_file, **kwargs): + super().__init__(**kwargs) - if pos_vocab is None: - raise ValueError("`pos_vocab` should be required.") + with open(pos_vocab_file, "r", encoding="utf-8") as f: + self.pos_vocab = json.load(f) - self.pos_vocab = pos_vocab - - # set pos_tagger + # set pos tagger tagger_cls = PosTaggerMap.get(self.lang, None) self.pos_tagger = tagger_cls() - def _tokenize(self, text: str) -> Tuple[List[str], List[str]]: + def _tokenize(self, text: str, **kwargs) -> Tuple[List[str], List[str]]: text = text.strip() pairs = self.pos_tagger.pos(text, return_surface=True) - tokens = ["", "▃"] + [pair[0] if pair[0] != " " else "▃" for pair in pairs] + tokens = ["", self.replacement] + [ + pair[0] if pair[0] != " " else self.replacement for pair in pairs + ] tags = [ pair[1] if pair[0] != " " else pairs[i + 1][1] for i, pair in enumerate(pairs) @@ -282,3 +283,110 @@ def _tokenize(self, text: str) -> Tuple[List[str], List[str]]: res_tags.append(tag) return tokens, res_tags + + @staticmethod + def _sanitize_kwargs( + method: Callable, **kwargs + ) -> Tuple[Dict[str, Any], Dict[str, Any]]: + orig_kwargs = {} + for arg_name in inspect.getfullargspec(method).args: + val = kwargs.pop(arg_name, None) + if val is not None: + orig_kwargs[arg_name] = val + return orig_kwargs, kwargs + + def convert_tags_to_ids( + self, tokens: Union[str, List[str]] + ) -> Union[int, List[int]]: + if tokens is None: + return None + + def _convert_tag_to_id(token: str) -> int: + if token is None: + return None + return self.pos_vocab.get(token, self.unk_token_id) + + if isinstance(tokens, str): + return _convert_tag_to_id(tokens) + + ids = [] + for token in tokens: + ids.append(_convert_tag_to_id(token)) + return ids + + def _get_input_ids(self, text: str, **kwargs) -> Tuple[List[int], List[int]]: + tokens, tags = self.tokenize(text, **kwargs) + return ( + self.convert_tokens_to_ids(tokens), + self.convert_tags_to_ids(tags), + ) + + def _encode_plus(self, text: str, text_pair: str, **kwargs) -> BatchEncoding: + if kwargs.pop("return_offsets_mapping", False): + raise NotImplementedError( + "return_offset_mapping is not available when using Python tokenizers. " + "To use this feature, change your tokenizer to one deriving from " + "transformers.PreTrainedTokenizerFast. " + "More information on available tokenizers at " + "https://github.com/huggingface/transformers/pull/2674" + ) + if kwargs.pop("is_split_into_words", False): + raise NotImplementedError( + "is_split_into_words is not available when using Pos Tokenizer. " + ) + orig_kwargs, kwargs = self._sanitize_kwargs(super()._encode_plus, **kwargs) + first_ids, first_tag_ids = self._get_input_ids(text, **kwargs) + second_ids = second_tag_ids = None + if text_pair is not None: + second_ids, second_tag_ids = self._get_input_ids(text_pair, **kwargs) + + batch_outputs = self.prepare_for_model( + ids=first_ids, pair_ids=second_ids, prepend_batch_axis=True, **orig_kwargs + ) + pos_outputs = self.prepare_for_model( + ids=first_tag_ids, + pair_ids=second_tag_ids, + prepend_batch_axis=True, + **orig_kwargs, + ) + batch_outputs.update({"segment_labels": pos_outputs["input_ids"]}) + return batch_outputs + + def _batch_encode_plus(self, batch_text_or_text_pairs, **kwargs) -> BatchEncoding: + if kwargs.pop("return_offsets_mapping", False): + raise NotImplementedError( + "return_offset_mapping is not available when using Python tokenizers. " + "To use this feature, change your tokenizer to one deriving from " + "transformers.PreTrainedTokenizerFast." + ) + if kwargs.pop("is_split_into_words", False): + raise NotImplementedError( + "is_split_into_words is not available when using Pos Tokenizer. " + ) + orig_kwargs, kwargs = self._sanitize_kwargs( + super()._batch_encode_plus, **kwargs + ) + input_ids = [] + tag_ids = [] + for ids_or_pair_ids in batch_text_or_text_pairs: + if not isinstance(ids_or_pair_ids, (list, tuple)): + ids, pair_ids = ids_or_pair_ids, None + else: + ids, pair_ids = ids_or_pair_ids + + first_ids, first_tag_ids = self._get_input_ids(ids, **kwargs) + second_ids = second_tag_ids = None + if pair_ids is not None: + second_ids, second_tag_ids = self._get_input_ids(pair_ids, **kwargs) + input_ids.append((first_ids, second_ids)) + tag_ids.append((first_tag_ids, second_tag_ids)) + + batch_outputs = self._batch_prepare_for_model(input_ids, **orig_kwargs) + pos_outputs = self._batch_prepare_for_model(tag_ids, **orig_kwargs) + batch_outputs.update({"segment_labels": pos_outputs["input_ids"]}) + return BatchEncoding(batch_outputs) + + def convert_tokens_to_string(self, tokens: List[str]) -> str: + """Converts a sequence of tokens (string) in a single string.""" + text = "".join(tokens) + return text.replace(self.replacement, " ").strip() diff --git a/dooly/utils/__init__.py b/dooly/utils/__init__.py new file mode 100644 index 0000000..01f41eb --- /dev/null +++ b/dooly/utils/__init__.py @@ -0,0 +1,16 @@ +from .hub import ( # noqa # pylint: disable=unused-import + register_subfolder, + recover_original_hf_bucket_url, + download_from_hf_hub, + DOOLY_HUB_NAME, +) + +from .import_utils import ( # noqa # pylint: disable=unused-import + is_available_mecab, + is_available_ipadic, + is_available_fugashi, + is_available_jieba, + is_available_nltk, + is_available_kss, +) +from .import_utils import _locate # noqa diff --git a/dooly/utils/hub.py b/dooly/utils/hub.py new file mode 100644 index 0000000..71a752d --- /dev/null +++ b/dooly/utils/hub.py @@ -0,0 +1,64 @@ +import os +import contextlib +from packaging import version +from typing import Callable, Any + +import transformers + + +DEFAULT_HUB_NAME = "jinmang2/dooly-hub" +DOOLY_HUB_NAME = os.environ.get("DOOLY_HUB_NAME", DEFAULT_HUB_NAME) + +hub_utils = [transformers.file_utils] +if version.parse(transformers.__version__) >= version.parse("4.16.0"): + hub_utils += [transformers.utils.hub] + +HF_CO_PREFIX = hub_utils[0].HUGGINGFACE_CO_PREFIX + + +def register_subfolder(func: Callable) -> Callable: + BASE_PREFIX = "https://huggingface.co/{model_id}/resolve/{revision}/" + + def wrapper(*args, **kwargs) -> Any: + subfolder = kwargs.pop("subfolder", None) + pretrained_model_name_or_path = kwargs.pop( + "pretrained_model_name_or_path", DOOLY_HUB_NAME + ) + _orig_hf_co_prefixes = [] + if pretrained_model_name_or_path == DOOLY_HUB_NAME: + for hub_util in hub_utils: + _orig_hf_co_prefixes.append(hub_util.HUGGINGFACE_CO_PREFIX) + hub_util.HUGGINGFACE_CO_PREFIX = BASE_PREFIX + hub_util.HUGGINGFACE_CO_PREFIX += f"{subfolder}/" if subfolder else "" + hub_util.HUGGINGFACE_CO_PREFIX += "{filename}" + kwargs.update( + dict( + subfolder=subfolder, + pretrained_model_name_or_path=pretrained_model_name_or_path, + ) + ) + output = func(*args, **kwargs) + for i, hf_co_prefix in enumerate(_orig_hf_co_prefixes): + hub_utils[i].HUGGINGFACE_CO_PREFIX = hf_co_prefix + return output + + return wrapper + + +@contextlib.contextmanager +def recover_original_hf_bucket_url(): + bucket_url_with_subfolder = hub_utils[0].HUGGINGFACE_CO_PREFIX + for hub_util in hub_utils: + hub_util.HUGGINGFACE_CO_PREFIX = HF_CO_PREFIX + yield + for hub_util in hub_utils: + hub_util.HUGGINGFACE_CO_PREFIX = bucket_url_with_subfolder + + +def download_from_hf_hub(filename: str, subfolder: str, hub_name: str = None) -> str: + hub_name = hub_name or DOOLY_HUB_NAME + hf_co_resolved_file = hub_utils[0].hf_bucket_url( + model_id=hub_name, filename=filename, subfolder=subfolder + ) + resolved_file_path = hub_utils[0].cached_path(hf_co_resolved_file) + return resolved_file_path diff --git a/dooly/utils/import_utils.py b/dooly/utils/import_utils.py new file mode 100644 index 0000000..f116c5b --- /dev/null +++ b/dooly/utils/import_utils.py @@ -0,0 +1,88 @@ +# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved +# Copyright 2022 The HuggingFace Team. All rights reserved. +import os +import importlib +from typing import Any + + +def is_available_mecab() -> bool: + _mecab = importlib.util.find_spec("mecab") + if os.name != "nt": + return _mecab is not None + else: + _eunjeon = importlib.util.find_spec("eunjeon") + return _mecab is not None or _eunjeon is not None + + +def is_available_ipadic() -> bool: + return importlib.util.find_spec("ipadic") + + +def is_available_fugashi() -> bool: + return importlib.util.find_spec("fugashi") + + +def is_available_jieba() -> bool: + return importlib.util.find_spec("jieba") + + +def is_available_nltk() -> bool: + return importlib.util.find_spec("nltk") + + +def is_available_kss() -> bool: + return importlib.util.find_spec("kss") + + +def _locate(path: str) -> Any: + """ + Locate an object by name or dotted path, importing as necessary. + This is similar to the pydoc function `locate`, except that it checks for + the module from the given path from back to front. + """ + if path == "": + raise ImportError("Empty path") + from importlib import import_module + from types import ModuleType + + parts = [part for part in path.split(".")] + for part in parts: + if not len(part): + raise ValueError( + f"Error loading '{path}': invalid dotstring." + + "\nRelative imports are not supported." + ) + assert len(parts) > 0 + part0 = parts[0] + try: + obj = import_module(part0) + except Exception as exc_import: + raise ImportError( + f"Error loading '{path}':\n{repr(exc_import)}" + + f"\nAre you sure that module '{part0}' is installed?" + ) from exc_import + for m in range(1, len(parts)): + part = parts[m] + try: + obj = getattr(obj, part) + except AttributeError as exc_attr: + parent_dotpath = ".".join(parts[:m]) + if isinstance(obj, ModuleType): + mod = ".".join(parts[: m + 1]) + try: + obj = import_module(mod) + continue + except ModuleNotFoundError as exc_import: + raise ImportError( + f"Error loading '{path}':\n{repr(exc_import)}" + + f"\nAre you sure that '{part}' is importable from module '{parent_dotpath}'?" + ) from exc_import + except Exception as exc_import: + raise ImportError( + f"Error loading '{path}':\n{repr(exc_import)}" + ) from exc_import + raise ImportError( + f"Error loading '{path}':\n{repr(exc_attr)}" + + f"\nAre you sure that '{part}' is an attribute of '{parent_dotpath}'?" + ) from exc_attr + return obj